Skip to main content

summa_core/segment/
vector_data.rs

1//! Vector index data structures shared between builder and reader
2
3use std::io;
4use std::mem::size_of;
5
6use crate::directories::{FileHandle, OwnedBytes};
7use crate::dsl::DenseVectorQuantization;
8use crate::segment::format::{DOC_ID_ENTRY_SIZE, FLAT_BINARY_HEADER_SIZE, FLAT_BINARY_MAGIC};
9use crate::structures::simd::{batch_f32_to_f16, batch_f32_to_u8, f16_to_f32, u8_to_f32};
10
11/// Dequantize raw bytes to f32 based on storage quantization.
12///
13/// `raw` is the quantized byte slice, `out` receives the f32 values.
14/// `num_floats` is the number of f32 values to produce (= num_vectors × dim).
15/// Data-first file layout guarantees alignment for f32/f16 access.
16#[inline]
17pub fn dequantize_raw(
18    raw: &[u8],
19    quant: DenseVectorQuantization,
20    num_floats: usize,
21    out: &mut [f32],
22) -> io::Result<()> {
23    if out.len() < num_floats {
24        return Err(io::Error::new(
25            io::ErrorKind::InvalidInput,
26            format!(
27                "dequantization output is too short: need {num_floats} floats, got {}",
28                out.len()
29            ),
30        ));
31    }
32
33    let element_size = match quant {
34        DenseVectorQuantization::F32 => size_of::<f32>(),
35        DenseVectorQuantization::F16 => size_of::<u16>(),
36        DenseVectorQuantization::UInt8 => size_of::<u8>(),
37        DenseVectorQuantization::Binary => {
38            return Err(io::Error::new(
39                io::ErrorKind::InvalidInput,
40                "binary vectors cannot be dequantized to f32",
41            ));
42        }
43    };
44    let expected_bytes = num_floats.checked_mul(element_size).ok_or_else(|| {
45        io::Error::new(
46            io::ErrorKind::InvalidInput,
47            "dequantization byte length overflows usize",
48        )
49    })?;
50    if raw.len() != expected_bytes {
51        return Err(io::Error::new(
52            io::ErrorKind::InvalidData,
53            format!(
54                "dequantization input length mismatch: need {expected_bytes} bytes, got {}",
55                raw.len()
56            ),
57        ));
58    }
59
60    match quant {
61        DenseVectorQuantization::F32 => {
62            if expected_bytes > 0
63                && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>())
64            {
65                return Err(io::Error::new(
66                    io::ErrorKind::InvalidData,
67                    "f32 vector data is not 4-byte aligned",
68                ));
69            }
70            out[..num_floats].copy_from_slice(unsafe {
71                // Safety: the exact byte length and f32 alignment were checked above.
72                std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats)
73            });
74        }
75        DenseVectorQuantization::F16 => {
76            if expected_bytes > 0
77                && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
78            {
79                return Err(io::Error::new(
80                    io::ErrorKind::InvalidData,
81                    "f16 vector data is not 2-byte aligned",
82                ));
83            }
84            let f16_slice = unsafe {
85                // Safety: the exact byte length and u16 alignment were checked above.
86                std::slice::from_raw_parts(raw.as_ptr() as *const u16, num_floats)
87            };
88            for (i, &h) in f16_slice.iter().enumerate() {
89                out[i] = f16_to_f32(h);
90            }
91        }
92        DenseVectorQuantization::UInt8 => {
93            for (i, &b) in raw.iter().enumerate() {
94                out[i] = u8_to_f32(b);
95            }
96        }
97        DenseVectorQuantization::Binary => unreachable!("validated above"),
98    }
99    Ok(())
100}
101
102/// Flat vector binary format helpers for writing.
103///
104/// Binary format v3:
105/// ```text
106/// [magic(u32)][dim(u32)][num_vectors(u32)][quant_type(u8)][padding(3)]
107/// [vectors: N×dim×element_size]
108/// [doc_ids: N×(u32+u16)]
109/// ```
110///
111/// `element_size` is determined by `quant_type`: f32=4, f16=2, uint8=1.
112/// Packed binary vectors instead use `dim / 8` bytes per row.
113/// [`LazyFlatVectorData`] retains a zero-copy view of the packed document/ordinal
114/// map and accesses vector payloads lazily via range reads (mmap on native files).
115pub struct FlatVectorData;
116
117impl FlatVectorData {
118    fn validate_shape(
119        dim: usize,
120        num_vectors: usize,
121        quant: DenseVectorQuantization,
122    ) -> io::Result<usize> {
123        if dim == 0 {
124            return Err(io::Error::new(
125                io::ErrorKind::InvalidInput,
126                "flat vector dimension must be greater than zero",
127            ));
128        }
129        if quant == DenseVectorQuantization::Binary && !dim.is_multiple_of(8) {
130            return Err(io::Error::new(
131                io::ErrorKind::InvalidInput,
132                format!("binary flat vector dimension must be a multiple of 8, got {dim}"),
133            ));
134        }
135        u32::try_from(dim).map_err(|_| {
136            io::Error::new(
137                io::ErrorKind::InvalidInput,
138                format!("flat vector dimension {dim} exceeds u32::MAX"),
139            )
140        })?;
141        u32::try_from(num_vectors).map_err(|_| {
142            io::Error::new(
143                io::ErrorKind::InvalidInput,
144                format!("flat vector count {num_vectors} exceeds u32::MAX"),
145            )
146        })?;
147
148        match quant {
149            DenseVectorQuantization::Binary => dim.checked_add(7).map(|bits| bits / 8),
150            _ => dim.checked_mul(quant.element_size()),
151        }
152        .ok_or_else(|| {
153            io::Error::new(
154                io::ErrorKind::InvalidInput,
155                "flat vector byte size overflows usize",
156            )
157        })
158    }
159
160    fn validate_doc_ids(doc_ids: &[(u32, u16)]) -> io::Result<()> {
161        if let Some(pair) = doc_ids.windows(2).find(|pair| pair[0] >= pair[1]) {
162            return Err(io::Error::new(
163                io::ErrorKind::InvalidInput,
164                format!(
165                    "flat vector doc map must be strictly sorted by (doc_id, ordinal), found {:?} before {:?}",
166                    pair[0], pair[1]
167                ),
168            ));
169        }
170        Ok(())
171    }
172
173    /// Validate a dense writer input completely before any bytes are emitted.
174    /// Returns the exact serialized size on success.
175    pub(crate) fn validate_dense_input(
176        dim: usize,
177        flat_vectors: &[f32],
178        doc_ids: &[(u32, u16)],
179        quant: DenseVectorQuantization,
180    ) -> io::Result<usize> {
181        if quant == DenseVectorQuantization::Binary {
182            return Err(io::Error::new(
183                io::ErrorKind::InvalidInput,
184                "binary quantization must use serialize_binary_from_bits_streaming",
185            ));
186        }
187        let num_vectors = doc_ids.len();
188        let expected_floats = num_vectors.checked_mul(dim).ok_or_else(|| {
189            io::Error::new(
190                io::ErrorKind::InvalidInput,
191                "flat f32 vector count overflows usize",
192            )
193        })?;
194        if flat_vectors.len() != expected_floats {
195            return Err(io::Error::new(
196                io::ErrorKind::InvalidInput,
197                format!(
198                    "flat vector input has {} floats, expected {num_vectors} x {dim} = {expected_floats}",
199                    flat_vectors.len()
200                ),
201            ));
202        }
203        Self::validate_doc_ids(doc_ids)?;
204        Self::serialized_binary_size(dim, num_vectors, quant)
205    }
206
207    /// Validate a packed-binary writer input completely before any bytes are
208    /// emitted. Returns the exact serialized size on success.
209    pub(crate) fn validate_binary_input(
210        dim_bits: usize,
211        packed_vectors: &[u8],
212        doc_ids: &[(u32, u16)],
213    ) -> io::Result<usize> {
214        let num_vectors = doc_ids.len();
215        let byte_len =
216            Self::validate_shape(dim_bits, num_vectors, DenseVectorQuantization::Binary)?;
217        let expected_bytes = num_vectors.checked_mul(byte_len).ok_or_else(|| {
218            io::Error::new(
219                io::ErrorKind::InvalidInput,
220                "packed binary vector size overflows usize",
221            )
222        })?;
223        if packed_vectors.len() != expected_bytes {
224            return Err(io::Error::new(
225                io::ErrorKind::InvalidInput,
226                format!(
227                    "packed binary input has {} bytes, expected {num_vectors} x {byte_len} = {expected_bytes}",
228                    packed_vectors.len()
229                ),
230            ));
231        }
232        Self::validate_doc_ids(doc_ids)?;
233        Self::serialized_binary_size(dim_bits, num_vectors, DenseVectorQuantization::Binary)
234    }
235
236    /// Write the binary header to a writer.
237    pub fn write_binary_header(
238        dim: usize,
239        num_vectors: usize,
240        quant: DenseVectorQuantization,
241        writer: &mut dyn std::io::Write,
242    ) -> std::io::Result<()> {
243        Self::validate_shape(dim, num_vectors, quant)?;
244        let dim = u32::try_from(dim).map_err(|_| {
245            io::Error::new(
246                io::ErrorKind::InvalidInput,
247                "flat vector dimension exceeds u32",
248            )
249        })?;
250        let num_vectors = u32::try_from(num_vectors).map_err(|_| {
251            io::Error::new(io::ErrorKind::InvalidInput, "flat vector count exceeds u32")
252        })?;
253        writer.write_all(&FLAT_BINARY_MAGIC.to_le_bytes())?;
254        writer.write_all(&dim.to_le_bytes())?;
255        writer.write_all(&num_vectors.to_le_bytes())?;
256        writer.write_all(&[quant.tag(), 0, 0, 0])?; // quant_type + 3 bytes padding
257        Ok(())
258    }
259
260    /// Compute the serialized size without actually serializing.
261    pub fn serialized_binary_size(
262        dim: usize,
263        num_vectors: usize,
264        quant: DenseVectorQuantization,
265    ) -> io::Result<usize> {
266        let bytes_per_vector = Self::validate_shape(dim, num_vectors, quant)?;
267        let vector_bytes = num_vectors.checked_mul(bytes_per_vector).ok_or_else(|| {
268            io::Error::new(
269                io::ErrorKind::InvalidInput,
270                "flat vector payload size overflows usize",
271            )
272        })?;
273        let doc_id_bytes = num_vectors.checked_mul(DOC_ID_ENTRY_SIZE).ok_or_else(|| {
274            io::Error::new(
275                io::ErrorKind::InvalidInput,
276                "flat vector doc-map size overflows usize",
277            )
278        })?;
279        FLAT_BINARY_HEADER_SIZE
280            .checked_add(vector_bytes)
281            .and_then(|size| size.checked_add(doc_id_bytes))
282            .ok_or_else(|| {
283                io::Error::new(
284                    io::ErrorKind::InvalidInput,
285                    "flat vector serialized size overflows usize",
286                )
287            })
288    }
289
290    /// Stream from flat f32 storage to a writer, quantizing on write.
291    ///
292    /// `flat_vectors` is contiguous storage of dim*n f32 floats.
293    /// Vectors are quantized to the specified format before writing.
294    pub fn serialize_binary_from_flat_streaming(
295        dim: usize,
296        flat_vectors: &[f32],
297        doc_ids: &[(u32, u16)],
298        quant: DenseVectorQuantization,
299        writer: &mut dyn std::io::Write,
300    ) -> std::io::Result<()> {
301        Self::validate_dense_input(dim, flat_vectors, doc_ids, quant)?;
302        let num_vectors = doc_ids.len();
303        Self::write_binary_header(dim, num_vectors, quant, writer)?;
304
305        match quant {
306            DenseVectorQuantization::F32 => {
307                let bytes: &[u8] = unsafe {
308                    std::slice::from_raw_parts(
309                        flat_vectors.as_ptr() as *const u8,
310                        std::mem::size_of_val(flat_vectors),
311                    )
312                };
313                writer.write_all(bytes)?;
314            }
315            DenseVectorQuantization::F16 => {
316                let mut buf = vec![0u16; dim];
317                for v in flat_vectors.chunks_exact(dim) {
318                    batch_f32_to_f16(v, &mut buf);
319                    let bytes: &[u8] =
320                        unsafe { std::slice::from_raw_parts(buf.as_ptr() as *const u8, dim * 2) };
321                    writer.write_all(bytes)?;
322                }
323            }
324            DenseVectorQuantization::UInt8 => {
325                let mut buf = vec![0u8; dim];
326                for v in flat_vectors.chunks_exact(dim) {
327                    batch_f32_to_u8(v, &mut buf);
328                    writer.write_all(&buf)?;
329                }
330            }
331            DenseVectorQuantization::Binary => unreachable!("validated above"),
332        }
333
334        for &(doc_id, ordinal) in doc_ids {
335            writer.write_all(&doc_id.to_le_bytes())?;
336            writer.write_all(&ordinal.to_le_bytes())?;
337        }
338
339        Ok(())
340    }
341
342    /// Stream packed binary vectors (pre-packed bytes) to a writer.
343    ///
344    /// `packed_vectors` is contiguous storage of num_vectors * byte_len bytes.
345    /// `dim_bits` is the number of bits (dimensions).
346    pub fn serialize_binary_from_bits_streaming(
347        dim_bits: usize,
348        packed_vectors: &[u8],
349        doc_ids: &[(u32, u16)],
350        writer: &mut dyn std::io::Write,
351    ) -> std::io::Result<()> {
352        Self::validate_binary_input(dim_bits, packed_vectors, doc_ids)?;
353        let num_vectors = doc_ids.len();
354
355        Self::write_binary_header(
356            dim_bits,
357            num_vectors,
358            DenseVectorQuantization::Binary,
359            writer,
360        )?;
361        writer.write_all(packed_vectors)?;
362
363        for &(doc_id, ordinal) in doc_ids {
364            writer.write_all(&doc_id.to_le_bytes())?;
365            writer.write_all(&ordinal.to_le_bytes())?;
366        }
367
368        Ok(())
369    }
370
371    /// Write raw pre-quantized vector bytes to a writer (for merger streaming).
372    ///
373    /// `raw_bytes` is already in the target quantized format.
374    pub fn write_raw_vector_bytes(
375        raw_bytes: &[u8],
376        writer: &mut dyn std::io::Write,
377    ) -> std::io::Result<()> {
378        writer.write_all(raw_bytes)
379    }
380}
381
382/// Lazy flat vector data — zero-copy doc_id index, vectors via range reads.
383///
384/// The doc_id index is kept as `OwnedBytes` (mmap-backed, zero heap copy).
385/// Exact vectors stay in flat storage or binary ANN code spans. Both layouts
386/// use the same document/ordinal lookup and lazy range-read interface.
387/// Element size depends on quantization: f32=4, f16=2, uint8=1 bytes/dim.
388///
389/// Used for:
390/// - Brute-force search (batched scoring with native-precision SIMD)
391/// - Reranking (read individual vectors by doc_id via binary search)
392/// - doc() hydration (dequantize to f32 for stored documents)
393/// - Merge streaming (chunked raw vector bytes + doc_id iteration)
394#[derive(Debug, Clone)]
395pub struct LazyFlatVectorData {
396    /// Vector dimension
397    pub dim: usize,
398    /// Total number of vectors
399    pub num_vectors: usize,
400    /// Number of distinct document IDs represented in the flat vector map.
401    num_docs_with_vectors: usize,
402    /// Storage quantization type
403    pub quantization: DenseVectorQuantization,
404    /// File-backed rows: document ID + ordinal, plus an address for ANN storage.
405    doc_ids_bytes: OwnedBytes,
406    /// Whether `doc_ids_bytes` holds the complete document map.
407    /// Validated once at open so per-vector lookups
408    /// need no length arithmetic; training-only readers leave it false.
409    has_doc_map: bool,
410    /// File handle for this field's flat or ANN region in the .vectors file
411    handle: FileHandle,
412    /// Byte offset within handle where raw vector data starts (after header)
413    vectors_offset: u64,
414    /// Bytes per vector in storage (cached: Binary = ceil(dim/8), else dim * element_size)
415    vbs: usize,
416    /// Exact byte length of the raw vector region, validated when opening.
417    vectors_byte_len: u64,
418    /// Exact codes live in ANN; doc-map entries contain span-relative addresses.
419    locations: Option<super::vector_locations::VectorLocations>,
420}
421
422impl LazyFlatVectorData {
423    /// Open from a lazy file slice pointing to the flat binary data region.
424    ///
425    /// Reads and validates the header and zero-copy document map. Vector data
426    /// stays lazy on disk.
427    pub async fn open(handle: FileHandle) -> io::Result<Self> {
428        Self::open_with_doc_limit(handle, None).await
429    }
430
431    /// Open flat vectors while also validating every referenced document ID.
432    ///
433    /// Segment readers pass their durable `num_docs` here. Keeping the public
434    /// `open` entry point is useful for standalone flat payloads and tests that
435    /// do not have segment metadata available.
436    pub(crate) async fn open_with_doc_limit(
437        handle: FileHandle,
438        total_docs: Option<u32>,
439    ) -> io::Result<Self> {
440        Self::open_impl(handle, total_docs, true).await
441    }
442
443    /// Open only the raw-vector region needed by global ANN training.
444    ///
445    /// The complete serialized shape is still checked, but the corpus-sized
446    /// document map is neither faulted in nor scanned: sampling addresses
447    /// vectors by global vector ordinal and never resolves document IDs.
448    pub(crate) async fn open_for_training(handle: FileHandle) -> io::Result<Self> {
449        Self::open_impl(handle, None, false).await
450    }
451
452    async fn open_impl(
453        handle: FileHandle,
454        total_docs: Option<u32>,
455        load_doc_map: bool,
456    ) -> io::Result<Self> {
457        let header_len = u64::try_from(FLAT_BINARY_HEADER_SIZE).map_err(|_| {
458            io::Error::new(
459                io::ErrorKind::InvalidData,
460                "flat vector header size does not fit in u64",
461            )
462        })?;
463        if handle.len() < header_len {
464            return Err(io::Error::new(
465                io::ErrorKind::UnexpectedEof,
466                format!(
467                    "flat vector payload is {} bytes, shorter than its {FLAT_BINARY_HEADER_SIZE}-byte header",
468                    handle.len()
469                ),
470            ));
471        }
472
473        // Read header: magic(4) + dim(4) + num_vectors(4) + quant_type(1) + pad(3) = 16 bytes
474        let header = handle.read_bytes_range(0..header_len).await?;
475        if header.len() != FLAT_BINARY_HEADER_SIZE {
476            return Err(io::Error::new(
477                io::ErrorKind::UnexpectedEof,
478                format!(
479                    "flat vector header read returned {} bytes, expected {FLAT_BINARY_HEADER_SIZE}",
480                    header.len()
481                ),
482            ));
483        }
484        let hdr = header.as_slice();
485
486        let magic = u32::from_le_bytes([hdr[0], hdr[1], hdr[2], hdr[3]]);
487        if magic != FLAT_BINARY_MAGIC {
488            return Err(io::Error::new(
489                io::ErrorKind::InvalidData,
490                "Invalid FlatVectorData binary magic",
491            ));
492        }
493
494        let dim = u32::from_le_bytes([hdr[4], hdr[5], hdr[6], hdr[7]]) as usize;
495        let num_vectors = u32::from_le_bytes([hdr[8], hdr[9], hdr[10], hdr[11]]) as usize;
496        let quantization = DenseVectorQuantization::from_tag(hdr[12]).ok_or_else(|| {
497            io::Error::new(
498                io::ErrorKind::InvalidData,
499                format!("Unknown quantization tag: {}", hdr[12]),
500            )
501        })?;
502        if hdr[13..] != [0, 0, 0] {
503            return Err(io::Error::new(
504                io::ErrorKind::InvalidData,
505                "flat vector header has non-zero reserved bytes",
506            ));
507        }
508
509        // Read doc_ids section as zero-copy OwnedBytes (6 bytes per vector)
510        let vbs =
511            FlatVectorData::validate_shape(dim, num_vectors, quantization).map_err(|error| {
512                io::Error::new(
513                    io::ErrorKind::InvalidData,
514                    format!("invalid flat vector shape: {error}"),
515                )
516            })?;
517        let vectors_byte_len_usize = num_vectors.checked_mul(vbs).ok_or_else(|| {
518            io::Error::new(
519                io::ErrorKind::InvalidData,
520                "flat vector payload size overflows usize",
521            )
522        })?;
523        let doc_ids_byte_len_usize =
524            num_vectors.checked_mul(DOC_ID_ENTRY_SIZE).ok_or_else(|| {
525                io::Error::new(
526                    io::ErrorKind::InvalidData,
527                    "flat vector doc-map size overflows usize",
528                )
529            })?;
530        let expected_len_usize = FLAT_BINARY_HEADER_SIZE
531            .checked_add(vectors_byte_len_usize)
532            .and_then(|size| size.checked_add(doc_ids_byte_len_usize))
533            .ok_or_else(|| {
534                io::Error::new(
535                    io::ErrorKind::InvalidData,
536                    "flat vector serialized size overflows usize",
537                )
538            })?;
539        let expected_len = u64::try_from(expected_len_usize).map_err(|_| {
540            io::Error::new(
541                io::ErrorKind::InvalidData,
542                "flat vector serialized size does not fit in u64",
543            )
544        })?;
545        if handle.len() != expected_len {
546            return Err(io::Error::new(
547                io::ErrorKind::InvalidData,
548                format!(
549                    "flat vector payload has {} bytes, expected exactly {expected_len}",
550                    handle.len()
551                ),
552            ));
553        }
554
555        let vectors_byte_len = u64::try_from(vectors_byte_len_usize).map_err(|_| {
556            io::Error::new(
557                io::ErrorKind::InvalidData,
558                "flat vector payload size does not fit in u64",
559            )
560        })?;
561        let doc_ids_byte_len = u64::try_from(doc_ids_byte_len_usize).map_err(|_| {
562            io::Error::new(
563                io::ErrorKind::InvalidData,
564                "flat vector doc-map size does not fit in u64",
565            )
566        })?;
567        let doc_ids_start = header_len.checked_add(vectors_byte_len).ok_or_else(|| {
568            io::Error::new(
569                io::ErrorKind::InvalidData,
570                "flat vector doc-map offset overflows u64",
571            )
572        })?;
573        let doc_ids_end = doc_ids_start.checked_add(doc_ids_byte_len).ok_or_else(|| {
574            io::Error::new(
575                io::ErrorKind::InvalidData,
576                "flat vector doc-map range overflows u64",
577            )
578        })?;
579
580        let doc_ids_bytes = if load_doc_map {
581            let bytes = handle.read_bytes_range(doc_ids_start..doc_ids_end).await?;
582            if bytes.len() != doc_ids_byte_len_usize {
583                return Err(io::Error::new(
584                    io::ErrorKind::UnexpectedEof,
585                    format!(
586                        "flat vector doc-map read returned {} bytes, expected {doc_ids_byte_len_usize}",
587                        bytes.len()
588                    ),
589                ));
590            }
591            bytes
592        } else {
593            OwnedBytes::empty()
594        };
595
596        let num_docs_with_vectors =
597            Self::validate_doc_map(&doc_ids_bytes, DOC_ID_ENTRY_SIZE, total_docs, |doc, _| {
598                Ok(doc)
599            })?;
600
601        debug_assert!(!load_doc_map || doc_ids_bytes.len() == doc_ids_byte_len_usize);
602        Ok(Self {
603            dim,
604            num_vectors,
605            num_docs_with_vectors,
606            quantization,
607            doc_ids_bytes,
608            has_doc_map: load_doc_map,
609            handle,
610            vectors_offset: header_len,
611            vbs,
612            vectors_byte_len,
613            locations: None,
614        })
615    }
616
617    fn validate_doc_map(
618        bytes: &OwnedBytes,
619        stride: usize,
620        total_docs: Option<u32>,
621        mut resolve: impl FnMut(u32, u16) -> io::Result<u32>,
622    ) -> io::Result<usize> {
623        let mut previous = None;
624        let mut num_docs_with_vectors = 0usize;
625        for (index, entry) in bytes.chunks_exact(stride).enumerate() {
626            // Admission walks this metadata sequentially. Keep at most two
627            // 64K-row windows ahead, never prefetch the vector-code corpus.
628            #[cfg(feature = "native")]
629            if index.is_multiple_of(64 * 1024) {
630                let start = index * stride;
631                let end = start.saturating_add(128 * 1024 * stride).min(bytes.len());
632                bytes.madvise_range(start..end, libc::MADV_WILLNEED);
633            }
634            #[cfg(not(feature = "native"))]
635            let _ = index;
636            let local_doc = u32::from_le_bytes([entry[0], entry[1], entry[2], entry[3]]);
637            let ordinal = u16::from_le_bytes([entry[4], entry[5]]);
638            let doc_id = resolve(local_doc, ordinal)?;
639            let current = (doc_id, ordinal);
640            if let Some(previous) = previous
641                && previous >= current
642            {
643                return Err(io::Error::new(
644                    io::ErrorKind::InvalidData,
645                    format!(
646                        "flat vector doc map must be strictly sorted by (doc_id, ordinal), found {previous:?} before {current:?}"
647                    ),
648                ));
649            }
650            if let Some(limit) = total_docs
651                && doc_id >= limit
652            {
653                return Err(io::Error::new(
654                    io::ErrorKind::InvalidData,
655                    format!(
656                        "flat vector doc map references document {doc_id}, but segment contains only {} documents",
657                        limit
658                    ),
659                ));
660            }
661            if previous.is_none_or(|(previous_doc_id, _)| previous_doc_id != doc_id) {
662                num_docs_with_vectors = num_docs_with_vectors.checked_add(1).ok_or_else(|| {
663                    io::Error::new(
664                        io::ErrorKind::InvalidData,
665                        "flat vector distinct-document count overflows usize",
666                    )
667                })?;
668            }
669            previous = Some(current);
670        }
671
672        Ok(num_docs_with_vectors)
673    }
674
675    /// Open a document lookup into this field's exact ANN payload. The
676    /// lookup is evictable metadata, and the code payload is never copied.
677    pub(crate) async fn open_indirect(
678        map: FileHandle,
679        ann: &crate::segment::ann_disk::AnnDiskIndex,
680        total_docs: u32,
681    ) -> io::Result<Self> {
682        use super::vector_locations::{ENTRY_SIZE, VectorLocations};
683        let bytes = map.read_bytes().await?;
684        if bytes.len() as u64 != map.len() {
685            return Err(io::Error::new(
686                io::ErrorKind::UnexpectedEof,
687                "truncated exact-vector lookup read",
688            ));
689        }
690        let opened = VectorLocations::open(bytes)?;
691        let handle = ann.exact_vectors_handle();
692        let dim = opened.dim;
693        let num_vectors = opened.count;
694        if opened.ann_len != handle.len() || dim != ann.header().dim {
695            return Err(io::Error::new(
696                io::ErrorKind::InvalidData,
697                "exact-vector lookup does not match its ANN payload",
698            ));
699        }
700        let vbs =
701            FlatVectorData::validate_shape(dim, num_vectors, DenseVectorQuantization::Binary)?;
702        let doc_ids_bytes = opened.rows;
703        let layout = opened.layout;
704        let num_docs_with_vectors = {
705            let spans = ann.exact_location_spans(layout.spans())?;
706            let mut addresses = layout.addresses(doc_ids_bytes.as_slice());
707            Self::validate_doc_map(
708                &doc_ids_bytes,
709                ENTRY_SIZE,
710                Some(total_docs),
711                |local_doc, ordinal| {
712                    let (base, span, row) =
713                        addresses.next().expect("validated lookup row count")?;
714                    let doc = local_doc.checked_add(base).ok_or_else(|| {
715                        io::Error::new(io::ErrorKind::InvalidData, "vector document base overflow")
716                    })?;
717                    spans[span].validate_label(row, doc, ordinal)?;
718                    Ok(doc)
719                },
720            )?
721        };
722        let vectors_byte_len = (num_vectors as u64)
723            .checked_mul(vbs as u64)
724            .ok_or_else(|| {
725                io::Error::new(
726                    io::ErrorKind::InvalidData,
727                    "exact-vector byte count overflow",
728                )
729            })?;
730        Ok(Self {
731            dim,
732            num_vectors,
733            num_docs_with_vectors,
734            quantization: DenseVectorQuantization::Binary,
735            doc_ids_bytes,
736            has_doc_map: true,
737            handle,
738            vectors_offset: 0,
739            vbs,
740            vectors_byte_len,
741            locations: Some(layout),
742        })
743    }
744
745    /// Whether exact codes are shared with the ANN payload.
746    pub fn is_ann_backed(&self) -> bool {
747        self.locations.is_some()
748    }
749
750    fn doc_map_stride(&self) -> usize {
751        if self.locations.is_some() {
752            super::vector_locations::ENTRY_SIZE
753        } else {
754            DOC_ID_ENTRY_SIZE
755        }
756    }
757
758    fn code_offset(&self, idx: usize) -> u64 {
759        let at = idx * super::vector_locations::ENTRY_SIZE + DOC_ID_ENTRY_SIZE;
760        let address = u64::from_le_bytes(self.doc_ids_bytes[at..at + 8].try_into().unwrap());
761        self.locations
762            .as_ref()
763            .expect("ANN-backed vector")
764            .code_offset(idx, address, self.vbs)
765            .expect("validated vector location")
766    }
767
768    #[cfg(feature = "native")]
769    pub(super) fn locations(&self) -> Option<&super::vector_locations::VectorLocations> {
770        self.locations.as_ref()
771    }
772
773    #[cfg(feature = "native")]
774    pub(super) fn location_rows(&self) -> &[u8] {
775        self.doc_ids_bytes.as_slice()
776    }
777
778    /// Longest physically contiguous run within a requested document range.
779    fn contiguous_rows(&self, start: usize, count: usize) -> usize {
780        if self.locations.is_none() || count < 2 {
781            return count;
782        }
783        let first = self.code_offset(start);
784        (1..count)
785            .find(|&i| self.code_offset(start + i) != first + i as u64 * self.vbs as u64)
786            .unwrap_or(count)
787    }
788
789    fn checked_vector_range(
790        &self,
791        start_idx: usize,
792        count: usize,
793    ) -> io::Result<(std::ops::Range<u64>, usize)> {
794        let end_idx = start_idx.checked_add(count).ok_or_else(|| {
795            io::Error::new(
796                io::ErrorKind::InvalidInput,
797                "flat vector index range overflows usize",
798            )
799        })?;
800        if end_idx > self.num_vectors {
801            return Err(io::Error::new(
802                io::ErrorKind::InvalidInput,
803                format!(
804                    "flat vector range {start_idx}..{end_idx} exceeds {} vectors",
805                    self.num_vectors
806                ),
807            ));
808        }
809
810        if self.locations.is_some() {
811            let len = count.checked_mul(self.vbs).ok_or_else(|| {
812                io::Error::new(io::ErrorKind::InvalidInput, "vector batch size overflow")
813            })?;
814            if count == 0 {
815                return Ok((0..0, 0));
816            }
817            if self.contiguous_rows(start_idx, count) != count {
818                return Err(io::Error::new(
819                    io::ErrorKind::InvalidInput,
820                    "vector batch is not physically contiguous",
821                ));
822            }
823            let start = self.code_offset(start_idx);
824            let end = start
825                .checked_add(len as u64)
826                .filter(|&end| end <= self.handle.len())
827                .ok_or_else(|| {
828                    io::Error::new(
829                        io::ErrorKind::InvalidData,
830                        "vector location outside ANN payload",
831                    )
832                })?;
833            return Ok((start..end, len));
834        }
835        let relative_offset = start_idx.checked_mul(self.vbs).ok_or_else(|| {
836            io::Error::new(
837                io::ErrorKind::InvalidData,
838                "flat vector byte offset overflows usize",
839            )
840        })?;
841        let byte_len = count.checked_mul(self.vbs).ok_or_else(|| {
842            io::Error::new(
843                io::ErrorKind::InvalidInput,
844                "flat vector byte length overflows usize",
845            )
846        })?;
847        let relative_offset = u64::try_from(relative_offset).map_err(|_| {
848            io::Error::new(
849                io::ErrorKind::InvalidData,
850                "flat vector byte offset does not fit in u64",
851            )
852        })?;
853        let byte_len_u64 = u64::try_from(byte_len).map_err(|_| {
854            io::Error::new(
855                io::ErrorKind::InvalidInput,
856                "flat vector byte length does not fit in u64",
857            )
858        })?;
859        let start = self
860            .vectors_offset
861            .checked_add(relative_offset)
862            .ok_or_else(|| {
863                io::Error::new(
864                    io::ErrorKind::InvalidData,
865                    "flat vector byte offset overflows u64",
866                )
867            })?;
868        let end = start.checked_add(byte_len_u64).ok_or_else(|| {
869            io::Error::new(
870                io::ErrorKind::InvalidData,
871                "flat vector byte range overflows u64",
872            )
873        })?;
874        let vectors_end = self
875            .vectors_offset
876            .checked_add(self.vectors_byte_len)
877            .ok_or_else(|| {
878                io::Error::new(
879                    io::ErrorKind::InvalidData,
880                    "flat vector payload boundary overflows u64",
881                )
882            })?;
883        if end > vectors_end || end > self.handle.len() {
884            return Err(io::Error::new(
885                io::ErrorKind::InvalidData,
886                format!(
887                    "flat vector byte range {start}..{end} exceeds payload boundary {vectors_end}"
888                ),
889            ));
890        }
891        Ok((start..end, byte_len))
892    }
893
894    /// Pin the doc-id map (priority 3: every rerank / top-k resolution
895    /// binary-searches it).
896    #[cfg(feature = "native")]
897    pub(crate) fn pin_doc_ids(
898        &mut self,
899        mode: crate::segment::pin::PinMode,
900        remaining: &mut u64,
901        report: &mut crate::segment::pin::PinReport,
902    ) {
903        crate::segment::pin::pin_section(
904            &mut self.doc_ids_bytes,
905            "flat doc_ids",
906            mode,
907            remaining,
908            report,
909        );
910    }
911
912    /// Advise the kernel that vector data will be accessed at random offsets.
913    ///
914    /// Disables kernel readahead for the raw vector region. Rerank reads
915    /// scattered ~vbs-sized records; default readahead pulls in 128KB per
916    /// fault, evicting useful pages in memory-bound environments.
917    /// No-op for non-mmap (RAM, HTTP) backing.
918    #[cfg(feature = "native")]
919    pub fn advise_random_access(&self) {
920        if self.locations.is_some() {
921            return;
922        }
923        let Some(vectors_end) = self.vectors_offset.checked_add(self.vectors_byte_len) else {
924            return;
925        };
926        self.handle
927            .madvise_range(self.vectors_offset..vectors_end, libc::MADV_RANDOM);
928    }
929
930    /// Prefetch the pages backing a sorted set of vector indexes (`MADV_WILLNEED`).
931    ///
932    /// Coalesces adjacent candidates into ranges so the kernel can overlap
933    /// the page-ins instead of taking one synchronous major fault per vector
934    /// during the rerank read loop. Indexes must be yielded in ascending order.
935    /// No-op for non-mmap backing.
936    #[cfg(feature = "native")]
937    pub fn prefetch_vectors(&self, sorted_flat_indexes: impl IntoIterator<Item = usize>) {
938        /// Gap (in bytes) below which two candidate ranges are merged into one advice call.
939        const COALESCE_GAP: u64 = 64 * 1024;
940        let mut ranges = sorted_flat_indexes.into_iter().filter_map(|idx| {
941            self.checked_vector_range(idx, 1)
942                .ok()
943                .map(|(range, _)| range)
944        });
945        let Some(first) = ranges.next() else {
946            return;
947        };
948        let mut run_start = first.start;
949        let mut run_end = first.end;
950        for range in ranges {
951            if range.start >= run_start && range.start <= run_end.saturating_add(COALESCE_GAP) {
952                run_end = run_end.max(range.end);
953            } else {
954                self.handle
955                    .madvise_range(run_start..run_end, libc::MADV_WILLNEED);
956                run_start = range.start;
957                run_end = range.end;
958            }
959        }
960        self.handle
961            .madvise_range(run_start..run_end, libc::MADV_WILLNEED);
962    }
963
964    /// Read a single vector by index, dequantized to f32.
965    ///
966    /// `out` must have length >= `self.dim`. Returns `Ok(())` on success.
967    /// Used for ANN training and doc() hydration where f32 is needed.
968    pub async fn read_vector_into(&self, idx: usize, out: &mut [f32]) -> io::Result<()> {
969        if out.len() < self.dim {
970            return Err(io::Error::new(
971                io::ErrorKind::InvalidInput,
972                format!(
973                    "flat vector output is too short: need {} floats, got {}",
974                    self.dim,
975                    out.len()
976                ),
977            ));
978        }
979        let bytes = self.read_vectors_batch(idx, 1).await?;
980        dequantize_raw(bytes.as_slice(), self.quantization, self.dim, out)
981    }
982
983    /// Read a single vector by index, dequantized to f32 (allocates a new `Vec<f32>`).
984    pub async fn get_vector(&self, idx: usize) -> io::Result<Vec<f32>> {
985        let mut vector = vec![0f32; self.dim];
986        self.read_vector_into(idx, &mut vector).await?;
987        Ok(vector)
988    }
989
990    /// Read a single vector's raw bytes (no dequantization) into a caller-provided buffer.
991    ///
992    /// `out` must have length >= `self.vector_byte_size()`.
993    /// Used for native-precision reranking where raw quantized bytes are scored directly.
994    pub async fn read_vector_raw_into(&self, idx: usize, out: &mut [u8]) -> io::Result<()> {
995        self.read_vector_prefix_raw_into(idx, self.vector_byte_size(), out)
996            .await
997    }
998
999    /// Read a prefix of one vector's raw bytes into a caller-provided buffer.
1000    ///
1001    /// This is used by Matryoshka scoring to avoid reading the unused tail of
1002    /// a vector. Unlike the old full-vector boundary, all caller-controlled
1003    /// sizes and offset arithmetic are checked in release builds.
1004    pub async fn read_vector_prefix_raw_into(
1005        &self,
1006        idx: usize,
1007        prefix_byte_len: usize,
1008        out: &mut [u8],
1009    ) -> io::Result<()> {
1010        let vbs = self.vector_byte_size();
1011        if prefix_byte_len > vbs {
1012            return Err(io::Error::new(
1013                io::ErrorKind::InvalidInput,
1014                format!(
1015                    "vector prefix is {prefix_byte_len} bytes, but a vector has only {vbs} bytes"
1016                ),
1017            ));
1018        }
1019        if out.len() < prefix_byte_len {
1020            return Err(io::Error::new(
1021                io::ErrorKind::InvalidInput,
1022                format!(
1023                    "vector prefix output is too short: need {prefix_byte_len} bytes, got {}",
1024                    out.len()
1025                ),
1026            ));
1027        }
1028        let (full_range, _) = self.checked_vector_range(idx, 1)?;
1029        if prefix_byte_len == 0 {
1030            return Ok(());
1031        }
1032        let prefix_byte_len_u64 = u64::try_from(prefix_byte_len).map_err(|_| {
1033            io::Error::new(
1034                io::ErrorKind::InvalidInput,
1035                "vector prefix length does not fit in u64",
1036            )
1037        })?;
1038        let byte_end = full_range
1039            .start
1040            .checked_add(prefix_byte_len_u64)
1041            .ok_or_else(|| {
1042                io::Error::new(
1043                    io::ErrorKind::InvalidData,
1044                    "vector byte range overflows u64",
1045                )
1046            })?;
1047        let bytes = self
1048            .handle
1049            .read_bytes_range(full_range.start..byte_end)
1050            .await?;
1051        if bytes.len() != prefix_byte_len {
1052            return Err(io::Error::new(
1053                io::ErrorKind::UnexpectedEof,
1054                format!(
1055                    "vector prefix read returned {} bytes, expected {prefix_byte_len}",
1056                    bytes.len()
1057                ),
1058            ));
1059        }
1060        out[..prefix_byte_len].copy_from_slice(bytes.as_slice());
1061        Ok(())
1062    }
1063
1064    /// Read a contiguous batch of raw quantized bytes by index range.
1065    ///
1066    /// Returns raw bytes for vectors `[start_idx..start_idx+count)`.
1067    /// Bytes are in native quantized format — pass to `batch_cosine_scores_f16/u8`
1068    /// or `batch_cosine_scores` (for f32) for scoring.
1069    pub async fn read_vectors_batch(
1070        &self,
1071        start_idx: usize,
1072        count: usize,
1073    ) -> io::Result<OwnedBytes> {
1074        if self.locations.is_some() {
1075            start_idx
1076                .checked_add(count)
1077                .filter(|&end| end <= self.num_vectors)
1078                .ok_or_else(|| {
1079                    io::Error::new(
1080                        io::ErrorKind::InvalidInput,
1081                        "vector batch outside logical range",
1082                    )
1083                })?;
1084            let len = count.checked_mul(self.vbs).ok_or_else(|| {
1085                io::Error::new(io::ErrorKind::InvalidInput, "vector batch size overflow")
1086            })?;
1087            if count == 0 {
1088                return Ok(OwnedBytes::empty());
1089            }
1090            let first_count = self.contiguous_rows(start_idx, count);
1091            if first_count < count {
1092                let bytes = self.handle.read_bytes().await?;
1093                return Ok(self.gather_vectors(bytes.as_slice(), start_idx, count, len));
1094            }
1095        }
1096        let (range, expected_len) = self.checked_vector_range(start_idx, count)?;
1097        let bytes = self.handle.read_bytes_range(range).await?;
1098        if bytes.len() != expected_len {
1099            return Err(io::Error::new(
1100                io::ErrorKind::UnexpectedEof,
1101                format!(
1102                    "flat vector batch read returned {} bytes, expected {expected_len}",
1103                    bytes.len()
1104                ),
1105            ));
1106        }
1107        Ok(bytes)
1108    }
1109
1110    /// Gather from the ANN reader's shared immutable byte owner. No per-row
1111    /// range-read callback, byte-owner clone, or second corpus copy is needed.
1112    fn gather_vectors(&self, bytes: &[u8], start: usize, count: usize, len: usize) -> OwnedBytes {
1113        let mut gathered = Vec::with_capacity(len);
1114        for row in start..start + count {
1115            let offset = self.code_offset(row) as usize;
1116            gathered.extend_from_slice(&bytes[offset..offset + self.vbs]);
1117        }
1118        OwnedBytes::new(gathered)
1119    }
1120
1121    /// Synchronous read of a single vector's raw bytes.
1122    #[cfg(feature = "sync")]
1123    pub fn read_vector_raw_into_sync(&self, idx: usize, out: &mut [u8]) -> io::Result<()> {
1124        let vbs = self.vector_byte_size();
1125        if out.len() < vbs {
1126            return Err(io::Error::new(
1127                io::ErrorKind::InvalidInput,
1128                format!(
1129                    "flat vector output is too short: need {vbs} bytes, got {}",
1130                    out.len()
1131                ),
1132            ));
1133        }
1134        let bytes = self.read_vectors_batch_sync(idx, 1)?;
1135        out[..vbs].copy_from_slice(bytes.as_slice());
1136        Ok(())
1137    }
1138
1139    /// Synchronous batch read of raw quantized bytes.
1140    #[cfg(feature = "sync")]
1141    pub fn read_vectors_batch_sync(
1142        &self,
1143        start_idx: usize,
1144        count: usize,
1145    ) -> io::Result<OwnedBytes> {
1146        if self.locations.is_some() {
1147            start_idx
1148                .checked_add(count)
1149                .filter(|&end| end <= self.num_vectors)
1150                .ok_or_else(|| {
1151                    io::Error::new(
1152                        io::ErrorKind::InvalidInput,
1153                        "vector batch outside logical range",
1154                    )
1155                })?;
1156            let len = count.checked_mul(self.vbs).ok_or_else(|| {
1157                io::Error::new(io::ErrorKind::InvalidInput, "vector batch size overflow")
1158            })?;
1159            if count == 0 {
1160                return Ok(OwnedBytes::empty());
1161            }
1162            let first_count = self.contiguous_rows(start_idx, count);
1163            if first_count < count {
1164                let bytes = self.handle.read_bytes_sync()?;
1165                return Ok(self.gather_vectors(bytes.as_slice(), start_idx, count, len));
1166            }
1167        }
1168        let (range, expected_len) = self.checked_vector_range(start_idx, count)?;
1169        let bytes = self.handle.read_bytes_range_sync(range)?;
1170        if bytes.len() != expected_len {
1171            return Err(io::Error::new(
1172                io::ErrorKind::UnexpectedEof,
1173                format!(
1174                    "flat vector batch read returned {} bytes, expected {expected_len}",
1175                    bytes.len()
1176                ),
1177            ));
1178        }
1179        Ok(bytes)
1180    }
1181
1182    /// Find flat index range for a given doc_id (non-allocating).
1183    ///
1184    /// Returns `(start_index, count)` — the flat vector index range for this doc_id.
1185    /// Use `get_doc_id(start + i)` for `i in 0..count` to read individual entries.
1186    /// More efficient than `flat_indexes_for_doc` as it avoids Vec allocation.
1187    pub fn flat_indexes_for_doc_range(&self, doc_id: u32) -> (usize, usize) {
1188        let n = self.num_vectors;
1189        let start = {
1190            let mut lo = 0usize;
1191            let mut hi = n;
1192            while lo < hi {
1193                let mid = lo + (hi - lo) / 2;
1194                if self.doc_id_at(mid) < doc_id {
1195                    lo = mid + 1;
1196                } else {
1197                    hi = mid;
1198                }
1199            }
1200            lo
1201        };
1202        let mut count = 0;
1203        let mut i = start;
1204        while i < n && self.doc_id_at(i) == doc_id {
1205            count += 1;
1206            i += 1;
1207        }
1208        (start, count)
1209    }
1210
1211    /// Find flat indexes for a given doc_id via binary search on sorted doc_ids.
1212    ///
1213    /// doc_ids are sorted by (doc_id, ordinal) — segment builder adds docs
1214    /// sequentially. Binary search runs directly on zero-copy mmap bytes.
1215    ///
1216    /// Returns `(start_index, entries)` where start_index is the flat vector index.
1217    pub fn flat_indexes_for_doc(&self, doc_id: u32) -> (usize, Vec<(u32, u16)>) {
1218        let n = self.num_vectors;
1219        // Binary search: find first entry where doc_id >= target
1220        let start = {
1221            let mut lo = 0usize;
1222            let mut hi = n;
1223            while lo < hi {
1224                let mid = lo + (hi - lo) / 2;
1225                if self.doc_id_at(mid) < doc_id {
1226                    lo = mid + 1;
1227                } else {
1228                    hi = mid;
1229                }
1230            }
1231            lo
1232        };
1233        // Collect entries with matching doc_id
1234        let mut entries = Vec::new();
1235        let mut i = start;
1236        while i < n {
1237            let (did, ord) = self.get_doc_id(i);
1238            if did != doc_id {
1239                break;
1240            }
1241            entries.push((did, ord));
1242            i += 1;
1243        }
1244        (start, entries)
1245    }
1246
1247    /// One packed doc-map entry. The map length was validated at open, so
1248    /// this is a flag test plus a single bounds check on the entry slice.
1249    #[inline]
1250    fn doc_map_entry(&self, idx: usize) -> &[u8; DOC_ID_ENTRY_SIZE] {
1251        assert!(
1252            self.has_doc_map,
1253            "document IDs are unavailable on a training-only flat-vector reader",
1254        );
1255        let off = idx * self.doc_map_stride();
1256        self.doc_ids_bytes[off..off + DOC_ID_ENTRY_SIZE]
1257            .try_into()
1258            .expect("doc-map entry slice is DOC_ID_ENTRY_SIZE bytes")
1259    }
1260
1261    /// Read doc_id at index from raw bytes (no ordinal).
1262    #[inline]
1263    fn doc_id_at(&self, idx: usize) -> u32 {
1264        let d = self.doc_map_entry(idx);
1265        u32::from_le_bytes([d[0], d[1], d[2], d[3]])
1266            + self
1267                .locations
1268                .as_ref()
1269                .map_or(0, |layout| layout.doc_base(idx))
1270    }
1271
1272    /// Get doc_id and ordinal at index (parsed from zero-copy mmap bytes).
1273    #[inline]
1274    pub fn get_doc_id(&self, idx: usize) -> (u32, u16) {
1275        let d = self.doc_map_entry(idx);
1276        let doc_id = u32::from_le_bytes([d[0], d[1], d[2], d[3]])
1277            + self
1278                .locations
1279                .as_ref()
1280                .map_or(0, |layout| layout.doc_base(idx));
1281        let ordinal = u16::from_le_bytes([d[4], d[5]]);
1282        (doc_id, ordinal)
1283    }
1284
1285    /// Bytes per vector in storage (cached).
1286    #[inline]
1287    pub fn vector_byte_size(&self) -> usize {
1288        self.vbs
1289    }
1290
1291    /// Number of distinct documents that have at least one vector in this field.
1292    #[inline]
1293    pub fn num_docs_with_vectors(&self) -> usize {
1294        self.num_docs_with_vectors
1295    }
1296
1297    /// Contiguous flat payload for bounded raw copying, including vectors
1298    /// larger than a copy window. ANN-backed document order is not contiguous.
1299    #[cfg(feature = "native")]
1300    pub(crate) fn flat_region(&self) -> Option<(&FileHandle, std::ops::Range<u64>)> {
1301        self.locations.is_none().then(|| {
1302            (
1303                &self.handle,
1304                self.vectors_offset..self.vectors_offset + self.vectors_byte_len,
1305            )
1306        })
1307    }
1308
1309    /// Logical exact-vector bytes; ANN-backed vectors share these bytes with ANN.
1310    pub fn vector_bytes_len(&self) -> u64 {
1311        self.vectors_byte_len
1312    }
1313
1314    /// Persisted lookup bytes, including its small directories, or zero for flat storage.
1315    pub fn exact_lookup_bytes(&self) -> u64 {
1316        self.locations
1317            .as_ref()
1318            .map_or(0, |layout| layout.serialized_len(self.num_vectors))
1319    }
1320
1321    /// Estimated heap usage — document IDs and vectors are file-backed.
1322    pub fn estimated_heap_bytes(&self) -> usize {
1323        size_of::<Self>()
1324            + self
1325                .locations
1326                .as_ref()
1327                .map_or(0, |layout| layout.heap_bytes())
1328    }
1329}
1330
1331#[cfg(test)]
1332mod tests {
1333    use super::*;
1334
1335    #[test]
1336    fn dequantize_raw_accepts_valid_storage_formats() {
1337        let f32_values = [1.25f32, -2.5];
1338        let f32_bytes = unsafe {
1339            // Safety: viewing an initialized f32 array as bytes is always valid.
1340            std::slice::from_raw_parts(
1341                f32_values.as_ptr().cast::<u8>(),
1342                std::mem::size_of_val(&f32_values),
1343            )
1344        };
1345        let mut out = [0.0; 2];
1346        dequantize_raw(f32_bytes, DenseVectorQuantization::F32, 2, &mut out).unwrap();
1347        assert_eq!(out, f32_values);
1348
1349        let f16_values = [0x3c00u16, 0xc000u16]; // 1.0, -2.0
1350        let f16_bytes = unsafe {
1351            // Safety: viewing an initialized u16 array as bytes is always valid.
1352            std::slice::from_raw_parts(
1353                f16_values.as_ptr().cast::<u8>(),
1354                std::mem::size_of_val(&f16_values),
1355            )
1356        };
1357        dequantize_raw(f16_bytes, DenseVectorQuantization::F16, 2, &mut out).unwrap();
1358        assert_eq!(out, [1.0, -2.0]);
1359
1360        dequantize_raw(&[0, u8::MAX], DenseVectorQuantization::UInt8, 2, &mut out).unwrap();
1361        assert_eq!(out, [u8_to_f32(0), u8_to_f32(u8::MAX)]);
1362    }
1363
1364    #[test]
1365    fn dequantize_raw_rejects_invalid_lengths_and_binary_storage() {
1366        let mut out = [0.0; 2];
1367        assert_eq!(
1368            dequantize_raw(&[0; 7], DenseVectorQuantization::F32, 2, &mut out)
1369                .unwrap_err()
1370                .kind(),
1371            io::ErrorKind::InvalidData
1372        );
1373        assert_eq!(
1374            dequantize_raw(&[0; 8], DenseVectorQuantization::F32, 2, &mut out[..1])
1375                .unwrap_err()
1376                .kind(),
1377            io::ErrorKind::InvalidInput
1378        );
1379        assert_eq!(
1380            dequantize_raw(&[], DenseVectorQuantization::Binary, 0, &mut [])
1381                .unwrap_err()
1382                .kind(),
1383            io::ErrorKind::InvalidInput
1384        );
1385    }
1386
1387    #[test]
1388    fn dequantize_raw_rejects_misaligned_typed_storage() {
1389        let storage = [0u8; 9];
1390        let offset = if (storage.as_ptr() as usize).is_multiple_of(4) {
1391            1
1392        } else {
1393            0
1394        };
1395        let raw = &storage[offset..offset + 8];
1396        assert!(!(raw.as_ptr() as usize).is_multiple_of(4));
1397
1398        let mut out = [0.0; 2];
1399        assert_eq!(
1400            dequantize_raw(raw, DenseVectorQuantization::F32, 2, &mut out)
1401                .unwrap_err()
1402                .kind(),
1403            io::ErrorKind::InvalidData
1404        );
1405    }
1406
1407    #[test]
1408    fn flat_vector_writers_reject_inconsistent_shapes_and_doc_maps() {
1409        let mut encoded = Vec::new();
1410        assert!(
1411            FlatVectorData::serialize_binary_from_flat_streaming(
1412                0,
1413                &[],
1414                &[],
1415                DenseVectorQuantization::F32,
1416                &mut encoded,
1417            )
1418            .is_err()
1419        );
1420        assert!(encoded.is_empty());
1421
1422        assert!(
1423            FlatVectorData::serialize_binary_from_flat_streaming(
1424                2,
1425                &[1.0],
1426                &[(0, 0)],
1427                DenseVectorQuantization::F32,
1428                &mut encoded,
1429            )
1430            .is_err()
1431        );
1432        assert!(encoded.is_empty());
1433
1434        assert!(
1435            FlatVectorData::serialize_binary_from_flat_streaming(
1436                1,
1437                &[1.0],
1438                &[(0, 0)],
1439                DenseVectorQuantization::Binary,
1440                &mut encoded,
1441            )
1442            .is_err()
1443        );
1444        assert!(encoded.is_empty());
1445
1446        assert!(
1447            FlatVectorData::serialize_binary_from_flat_streaming(
1448                1,
1449                &[1.0, 2.0],
1450                &[(1, 0), (0, 0)],
1451                DenseVectorQuantization::F32,
1452                &mut encoded,
1453            )
1454            .is_err()
1455        );
1456        assert!(encoded.is_empty());
1457
1458        assert!(
1459            FlatVectorData::serialize_binary_from_flat_streaming(
1460                1,
1461                &[1.0, 2.0],
1462                &[(0, 0), (0, 0)],
1463                DenseVectorQuantization::F32,
1464                &mut encoded,
1465            )
1466            .is_err()
1467        );
1468        assert!(encoded.is_empty());
1469
1470        assert!(
1471            FlatVectorData::serialize_binary_from_bits_streaming(7, &[0], &[(0, 0)], &mut encoded,)
1472                .is_err()
1473        );
1474        assert!(encoded.is_empty());
1475
1476        assert!(
1477            FlatVectorData::serialize_binary_from_bits_streaming(8, &[], &[(0, 0)], &mut encoded,)
1478                .is_err()
1479        );
1480        assert!(encoded.is_empty());
1481
1482        let vectors = [1.0f32, 2.0, 3.0, 4.0];
1483        let doc_ids = [(0, 0), (1, 0)];
1484        FlatVectorData::serialize_binary_from_flat_streaming(
1485            2,
1486            &vectors,
1487            &doc_ids,
1488            DenseVectorQuantization::F32,
1489            &mut encoded,
1490        )
1491        .unwrap();
1492        assert_eq!(
1493            encoded.len(),
1494            FlatVectorData::serialized_binary_size(2, 2, DenseVectorQuantization::F32).unwrap()
1495        );
1496    }
1497
1498    fn encoded_two_vector_payload() -> Vec<u8> {
1499        let mut encoded = Vec::new();
1500        FlatVectorData::serialize_binary_from_flat_streaming(
1501            2,
1502            &[1.0, 2.0, 3.0, 4.0],
1503            &[(0, 0), (1, 0)],
1504            DenseVectorQuantization::F32,
1505            &mut encoded,
1506        )
1507        .unwrap();
1508        encoded
1509    }
1510
1511    #[tokio::test]
1512    async fn flat_vector_open_rejects_corrupt_layout_and_doc_map() {
1513        let valid = encoded_two_vector_payload();
1514
1515        let mut multi_value = Vec::new();
1516        FlatVectorData::serialize_binary_from_flat_streaming(
1517            1,
1518            &[1.0, 2.0, 3.0],
1519            &[(0, 0), (0, 1), (2, 0)],
1520            DenseVectorQuantization::F32,
1521            &mut multi_value,
1522        )
1523        .unwrap();
1524        let multi_value = LazyFlatVectorData::open_with_doc_limit(
1525            FileHandle::from_bytes(OwnedBytes::new(multi_value)),
1526            Some(3),
1527        )
1528        .await
1529        .unwrap();
1530        assert_eq!(multi_value.num_docs_with_vectors(), 2);
1531
1532        let mut trailing = valid.clone();
1533        trailing.push(0);
1534        assert!(
1535            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(trailing)))
1536                .await
1537                .is_err()
1538        );
1539
1540        let mut truncated = valid.clone();
1541        truncated.pop();
1542        assert!(
1543            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(truncated)))
1544                .await
1545                .is_err()
1546        );
1547
1548        let mut reserved = valid.clone();
1549        reserved[13] = 1;
1550        assert!(
1551            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(reserved)))
1552                .await
1553                .is_err()
1554        );
1555
1556        let doc_map_start = FLAT_BINARY_HEADER_SIZE + 2 * 2 * size_of::<f32>();
1557        let mut unsorted = valid.clone();
1558        let (first, second) = unsorted[doc_map_start..doc_map_start + 2 * DOC_ID_ENTRY_SIZE]
1559            .split_at_mut(DOC_ID_ENTRY_SIZE);
1560        first.swap_with_slice(second);
1561        assert!(
1562            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(unsorted)))
1563                .await
1564                .is_err()
1565        );
1566
1567        let mut duplicate = valid.clone();
1568        duplicate.copy_within(
1569            doc_map_start..doc_map_start + DOC_ID_ENTRY_SIZE,
1570            doc_map_start + DOC_ID_ENTRY_SIZE,
1571        );
1572        assert!(
1573            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(duplicate)))
1574                .await
1575                .is_err()
1576        );
1577
1578        assert!(
1579            LazyFlatVectorData::open_with_doc_limit(
1580                FileHandle::from_bytes(OwnedBytes::new(valid)),
1581                Some(1),
1582            )
1583            .await
1584            .is_err()
1585        );
1586
1587        let mut invalid_binary = Vec::new();
1588        FlatVectorData::serialize_binary_from_bits_streaming(
1589            8,
1590            &[0],
1591            &[(0, 0)],
1592            &mut invalid_binary,
1593        )
1594        .unwrap();
1595        invalid_binary[4..8].copy_from_slice(&7u32.to_le_bytes());
1596        assert!(
1597            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(invalid_binary)))
1598                .await
1599                .is_err()
1600        );
1601    }
1602
1603    #[tokio::test]
1604    async fn flat_vector_batch_and_dequantized_reads_are_checked() {
1605        let flat = LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(
1606            encoded_two_vector_payload(),
1607        )))
1608        .await
1609        .unwrap();
1610
1611        assert_eq!(flat.read_vectors_batch(0, 2).await.unwrap().len(), 16);
1612        assert_eq!(flat.read_vectors_batch(2, 0).await.unwrap().len(), 0);
1613        assert!(flat.read_vectors_batch(1, 2).await.is_err());
1614        assert!(flat.read_vectors_batch(usize::MAX, 1).await.is_err());
1615        assert!(flat.read_vectors_batch(0, usize::MAX).await.is_err());
1616
1617        let mut values = [0.0; 2];
1618        flat.read_vector_into(1, &mut values).await.unwrap();
1619        assert_eq!(values, [3.0, 4.0]);
1620        assert!(flat.read_vector_into(2, &mut values).await.is_err());
1621        assert!(flat.read_vector_into(0, &mut values[..1]).await.is_err());
1622
1623        #[cfg(feature = "sync")]
1624        {
1625            assert_eq!(flat.read_vectors_batch_sync(0, 2).unwrap().len(), 16);
1626            assert!(flat.read_vectors_batch_sync(1, 2).is_err());
1627            assert!(flat.read_vectors_batch_sync(usize::MAX, 1).is_err());
1628            let mut too_short = [0; 7];
1629            assert!(flat.read_vector_raw_into_sync(0, &mut too_short).is_err());
1630        }
1631    }
1632
1633    #[cfg(not(target_arch = "wasm32"))]
1634    #[tokio::test]
1635    async fn flat_vector_reads_reject_short_lazy_range_results() {
1636        let payload = std::sync::Arc::new(encoded_two_vector_payload());
1637        let payload_len = payload.len() as u64;
1638        let read_payload = std::sync::Arc::clone(&payload);
1639        let read_fn: crate::directories::RangeReadFn = std::sync::Arc::new(move |range| {
1640            let payload = std::sync::Arc::clone(&read_payload);
1641            Box::pin(async move {
1642                let start = usize::try_from(range.start).unwrap();
1643                let mut end = usize::try_from(range.end).unwrap();
1644                // Header and doc-map reads are exact, allowing open to finish.
1645                // Raw vector reads deliberately violate the range-read contract.
1646                if range.start == FLAT_BINARY_HEADER_SIZE as u64 {
1647                    end -= 1;
1648                }
1649                Ok(OwnedBytes::new(payload[start..end].to_vec()))
1650            })
1651        });
1652        let flat = LazyFlatVectorData::open(FileHandle::lazy(payload_len, read_fn))
1653            .await
1654            .unwrap();
1655
1656        let error = flat.read_vectors_batch(0, 1).await.unwrap_err();
1657        assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
1658        let mut raw = [0; 8];
1659        let error = flat.read_vector_raw_into(0, &mut raw).await.unwrap_err();
1660        assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
1661    }
1662
1663    #[tokio::test]
1664    async fn vector_prefix_reads_are_checked_and_do_not_fetch_the_tail() {
1665        let vectors = [1.0f32, 2.0, 3.0, 4.0];
1666        let doc_ids = [(0, 0), (1, 0)];
1667        let mut encoded = Vec::new();
1668        FlatVectorData::serialize_binary_from_flat_streaming(
1669            2,
1670            &vectors,
1671            &doc_ids,
1672            DenseVectorQuantization::F32,
1673            &mut encoded,
1674        )
1675        .unwrap();
1676        let flat = LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(encoded)))
1677            .await
1678            .unwrap();
1679
1680        let mut prefix = [0xa5; 8];
1681        flat.read_vector_prefix_raw_into(1, 4, &mut prefix)
1682            .await
1683            .unwrap();
1684        assert_eq!(&prefix[..4], &3.0f32.to_ne_bytes());
1685        assert_eq!(&prefix[4..], &[0xa5; 4]);
1686
1687        let mut full = [0; 8];
1688        flat.read_vector_raw_into(1, &mut full).await.unwrap();
1689        assert_eq!(&full[..4], &3.0f32.to_ne_bytes());
1690        assert_eq!(&full[4..], &4.0f32.to_ne_bytes());
1691
1692        assert!(
1693            flat.read_vector_prefix_raw_into(2, 4, &mut prefix)
1694                .await
1695                .is_err()
1696        );
1697        assert!(
1698            flat.read_vector_prefix_raw_into(0, 9, &mut prefix)
1699                .await
1700                .is_err()
1701        );
1702        assert!(
1703            flat.read_vector_prefix_raw_into(0, 4, &mut prefix[..3])
1704                .await
1705                .is_err()
1706        );
1707    }
1708}