Skip to main content

hermes_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/// Reading is handled by [`LazyFlatVectorData`] which loads only doc_ids into memory
113/// and accesses vector data lazily via mmap-backed range reads.
114pub struct FlatVectorData;
115
116impl FlatVectorData {
117    fn validate_shape(
118        dim: usize,
119        num_vectors: usize,
120        quant: DenseVectorQuantization,
121    ) -> io::Result<usize> {
122        if dim == 0 {
123            return Err(io::Error::new(
124                io::ErrorKind::InvalidInput,
125                "flat vector dimension must be greater than zero",
126            ));
127        }
128        if quant == DenseVectorQuantization::Binary && !dim.is_multiple_of(8) {
129            return Err(io::Error::new(
130                io::ErrorKind::InvalidInput,
131                format!("binary flat vector dimension must be a multiple of 8, got {dim}"),
132            ));
133        }
134        u32::try_from(dim).map_err(|_| {
135            io::Error::new(
136                io::ErrorKind::InvalidInput,
137                format!("flat vector dimension {dim} exceeds u32::MAX"),
138            )
139        })?;
140        u32::try_from(num_vectors).map_err(|_| {
141            io::Error::new(
142                io::ErrorKind::InvalidInput,
143                format!("flat vector count {num_vectors} exceeds u32::MAX"),
144            )
145        })?;
146
147        match quant {
148            DenseVectorQuantization::Binary => dim.checked_add(7).map(|bits| bits / 8),
149            _ => dim.checked_mul(quant.element_size()),
150        }
151        .ok_or_else(|| {
152            io::Error::new(
153                io::ErrorKind::InvalidInput,
154                "flat vector byte size overflows usize",
155            )
156        })
157    }
158
159    fn validate_doc_ids(doc_ids: &[(u32, u16)]) -> io::Result<()> {
160        if let Some(pair) = doc_ids.windows(2).find(|pair| pair[0] >= pair[1]) {
161            return Err(io::Error::new(
162                io::ErrorKind::InvalidInput,
163                format!(
164                    "flat vector doc map must be strictly sorted by (doc_id, ordinal), found {:?} before {:?}",
165                    pair[0], pair[1]
166                ),
167            ));
168        }
169        Ok(())
170    }
171
172    /// Validate a dense writer input completely before any bytes are emitted.
173    /// Returns the exact serialized size on success.
174    pub(crate) fn validate_dense_input(
175        dim: usize,
176        flat_vectors: &[f32],
177        doc_ids: &[(u32, u16)],
178        quant: DenseVectorQuantization,
179    ) -> io::Result<usize> {
180        if quant == DenseVectorQuantization::Binary {
181            return Err(io::Error::new(
182                io::ErrorKind::InvalidInput,
183                "binary quantization must use serialize_binary_from_bits_streaming",
184            ));
185        }
186        let num_vectors = doc_ids.len();
187        let expected_floats = num_vectors.checked_mul(dim).ok_or_else(|| {
188            io::Error::new(
189                io::ErrorKind::InvalidInput,
190                "flat f32 vector count overflows usize",
191            )
192        })?;
193        if flat_vectors.len() != expected_floats {
194            return Err(io::Error::new(
195                io::ErrorKind::InvalidInput,
196                format!(
197                    "flat vector input has {} floats, expected {num_vectors} x {dim} = {expected_floats}",
198                    flat_vectors.len()
199                ),
200            ));
201        }
202        Self::validate_doc_ids(doc_ids)?;
203        Self::serialized_binary_size(dim, num_vectors, quant)
204    }
205
206    /// Validate a packed-binary writer input completely before any bytes are
207    /// emitted. Returns the exact serialized size on success.
208    pub(crate) fn validate_binary_input(
209        dim_bits: usize,
210        packed_vectors: &[u8],
211        doc_ids: &[(u32, u16)],
212    ) -> io::Result<usize> {
213        let num_vectors = doc_ids.len();
214        let byte_len =
215            Self::validate_shape(dim_bits, num_vectors, DenseVectorQuantization::Binary)?;
216        let expected_bytes = num_vectors.checked_mul(byte_len).ok_or_else(|| {
217            io::Error::new(
218                io::ErrorKind::InvalidInput,
219                "packed binary vector size overflows usize",
220            )
221        })?;
222        if packed_vectors.len() != expected_bytes {
223            return Err(io::Error::new(
224                io::ErrorKind::InvalidInput,
225                format!(
226                    "packed binary input has {} bytes, expected {num_vectors} x {byte_len} = {expected_bytes}",
227                    packed_vectors.len()
228                ),
229            ));
230        }
231        Self::validate_doc_ids(doc_ids)?;
232        Self::serialized_binary_size(dim_bits, num_vectors, DenseVectorQuantization::Binary)
233    }
234
235    /// Write the binary header to a writer.
236    pub fn write_binary_header(
237        dim: usize,
238        num_vectors: usize,
239        quant: DenseVectorQuantization,
240        writer: &mut dyn std::io::Write,
241    ) -> std::io::Result<()> {
242        Self::validate_shape(dim, num_vectors, quant)?;
243        let dim = u32::try_from(dim).map_err(|_| {
244            io::Error::new(
245                io::ErrorKind::InvalidInput,
246                "flat vector dimension exceeds u32",
247            )
248        })?;
249        let num_vectors = u32::try_from(num_vectors).map_err(|_| {
250            io::Error::new(io::ErrorKind::InvalidInput, "flat vector count exceeds u32")
251        })?;
252        writer.write_all(&FLAT_BINARY_MAGIC.to_le_bytes())?;
253        writer.write_all(&dim.to_le_bytes())?;
254        writer.write_all(&num_vectors.to_le_bytes())?;
255        writer.write_all(&[quant.tag(), 0, 0, 0])?; // quant_type + 3 bytes padding
256        Ok(())
257    }
258
259    /// Compute the serialized size without actually serializing.
260    pub fn serialized_binary_size(
261        dim: usize,
262        num_vectors: usize,
263        quant: DenseVectorQuantization,
264    ) -> io::Result<usize> {
265        let bytes_per_vector = Self::validate_shape(dim, num_vectors, quant)?;
266        let vector_bytes = num_vectors.checked_mul(bytes_per_vector).ok_or_else(|| {
267            io::Error::new(
268                io::ErrorKind::InvalidInput,
269                "flat vector payload size overflows usize",
270            )
271        })?;
272        let doc_id_bytes = num_vectors.checked_mul(DOC_ID_ENTRY_SIZE).ok_or_else(|| {
273            io::Error::new(
274                io::ErrorKind::InvalidInput,
275                "flat vector doc-map size overflows usize",
276            )
277        })?;
278        FLAT_BINARY_HEADER_SIZE
279            .checked_add(vector_bytes)
280            .and_then(|size| size.checked_add(doc_id_bytes))
281            .ok_or_else(|| {
282                io::Error::new(
283                    io::ErrorKind::InvalidInput,
284                    "flat vector serialized size overflows usize",
285                )
286            })
287    }
288
289    /// Stream from flat f32 storage to a writer, quantizing on write.
290    ///
291    /// `flat_vectors` is contiguous storage of dim*n f32 floats.
292    /// Vectors are quantized to the specified format before writing.
293    pub fn serialize_binary_from_flat_streaming(
294        dim: usize,
295        flat_vectors: &[f32],
296        doc_ids: &[(u32, u16)],
297        quant: DenseVectorQuantization,
298        writer: &mut dyn std::io::Write,
299    ) -> std::io::Result<()> {
300        Self::validate_dense_input(dim, flat_vectors, doc_ids, quant)?;
301        let num_vectors = doc_ids.len();
302        Self::write_binary_header(dim, num_vectors, quant, writer)?;
303
304        match quant {
305            DenseVectorQuantization::F32 => {
306                let bytes: &[u8] = unsafe {
307                    std::slice::from_raw_parts(
308                        flat_vectors.as_ptr() as *const u8,
309                        std::mem::size_of_val(flat_vectors),
310                    )
311                };
312                writer.write_all(bytes)?;
313            }
314            DenseVectorQuantization::F16 => {
315                let mut buf = vec![0u16; dim];
316                for v in flat_vectors.chunks_exact(dim) {
317                    batch_f32_to_f16(v, &mut buf);
318                    let bytes: &[u8] =
319                        unsafe { std::slice::from_raw_parts(buf.as_ptr() as *const u8, dim * 2) };
320                    writer.write_all(bytes)?;
321                }
322            }
323            DenseVectorQuantization::UInt8 => {
324                let mut buf = vec![0u8; dim];
325                for v in flat_vectors.chunks_exact(dim) {
326                    batch_f32_to_u8(v, &mut buf);
327                    writer.write_all(&buf)?;
328                }
329            }
330            DenseVectorQuantization::Binary => unreachable!("validated above"),
331        }
332
333        for &(doc_id, ordinal) in doc_ids {
334            writer.write_all(&doc_id.to_le_bytes())?;
335            writer.write_all(&ordinal.to_le_bytes())?;
336        }
337
338        Ok(())
339    }
340
341    /// Stream packed binary vectors (pre-packed bytes) to a writer.
342    ///
343    /// `packed_vectors` is contiguous storage of num_vectors * byte_len bytes.
344    /// `dim_bits` is the number of bits (dimensions).
345    pub fn serialize_binary_from_bits_streaming(
346        dim_bits: usize,
347        packed_vectors: &[u8],
348        doc_ids: &[(u32, u16)],
349        writer: &mut dyn std::io::Write,
350    ) -> std::io::Result<()> {
351        Self::validate_binary_input(dim_bits, packed_vectors, doc_ids)?;
352        let num_vectors = doc_ids.len();
353
354        Self::write_binary_header(
355            dim_bits,
356            num_vectors,
357            DenseVectorQuantization::Binary,
358            writer,
359        )?;
360        writer.write_all(packed_vectors)?;
361
362        for &(doc_id, ordinal) in doc_ids {
363            writer.write_all(&doc_id.to_le_bytes())?;
364            writer.write_all(&ordinal.to_le_bytes())?;
365        }
366
367        Ok(())
368    }
369
370    /// Write raw pre-quantized vector bytes to a writer (for merger streaming).
371    ///
372    /// `raw_bytes` is already in the target quantized format.
373    pub fn write_raw_vector_bytes(
374        raw_bytes: &[u8],
375        writer: &mut dyn std::io::Write,
376    ) -> std::io::Result<()> {
377        writer.write_all(raw_bytes)
378    }
379}
380
381/// Lazy flat vector data — zero-copy doc_id index, vectors via range reads.
382///
383/// The doc_id index is kept as `OwnedBytes` (mmap-backed, zero heap copy).
384/// Vector data stays on disk and is accessed via mmap-backed range reads.
385/// Element size depends on quantization: f32=4, f16=2, uint8=1 bytes/dim.
386///
387/// Used for:
388/// - Brute-force search (batched scoring with native-precision SIMD)
389/// - Reranking (read individual vectors by doc_id via binary search)
390/// - doc() hydration (dequantize to f32 for stored documents)
391/// - Merge streaming (chunked raw vector bytes + doc_id iteration)
392#[derive(Debug, Clone)]
393pub struct LazyFlatVectorData {
394    /// Vector dimension
395    pub dim: usize,
396    /// Total number of vectors
397    pub num_vectors: usize,
398    /// Number of distinct document IDs represented in the flat vector map.
399    num_docs_with_vectors: usize,
400    /// Storage quantization type
401    pub quantization: DenseVectorQuantization,
402    /// Zero-copy doc_id index: packed [u32_le doc_id + u16_le ordinal] × num_vectors
403    doc_ids_bytes: OwnedBytes,
404    /// File handle for this field's flat data region in the .vectors file
405    handle: FileHandle,
406    /// Byte offset within handle where raw vector data starts (after header)
407    vectors_offset: u64,
408    /// Bytes per vector in storage (cached: Binary = ceil(dim/8), else dim * element_size)
409    vbs: usize,
410    /// Exact byte length of the raw vector region, validated when opening.
411    vectors_byte_len: u64,
412}
413
414impl LazyFlatVectorData {
415    /// Open from a lazy file slice pointing to the flat binary data region.
416    ///
417    /// Reads and validates the header and zero-copy document map. Vector data
418    /// stays lazy on disk.
419    pub async fn open(handle: FileHandle) -> io::Result<Self> {
420        Self::open_with_doc_limit(handle, None).await
421    }
422
423    /// Open flat vectors while also validating every referenced document ID.
424    ///
425    /// Segment readers pass their durable `num_docs` here. Keeping the public
426    /// `open` entry point is useful for standalone flat payloads and tests that
427    /// do not have segment metadata available.
428    pub(crate) async fn open_with_doc_limit(
429        handle: FileHandle,
430        total_docs: Option<u32>,
431    ) -> io::Result<Self> {
432        Self::open_impl(handle, total_docs, true).await
433    }
434
435    /// Open only the raw-vector region needed by global ANN training.
436    ///
437    /// The complete serialized shape is still checked, but the corpus-sized
438    /// document map is neither faulted in nor scanned: sampling addresses
439    /// vectors by global vector ordinal and never resolves document IDs.
440    pub(crate) async fn open_for_training(handle: FileHandle) -> io::Result<Self> {
441        Self::open_impl(handle, None, false).await
442    }
443
444    async fn open_impl(
445        handle: FileHandle,
446        total_docs: Option<u32>,
447        load_doc_map: bool,
448    ) -> io::Result<Self> {
449        let header_len = u64::try_from(FLAT_BINARY_HEADER_SIZE).map_err(|_| {
450            io::Error::new(
451                io::ErrorKind::InvalidData,
452                "flat vector header size does not fit in u64",
453            )
454        })?;
455        if handle.len() < header_len {
456            return Err(io::Error::new(
457                io::ErrorKind::UnexpectedEof,
458                format!(
459                    "flat vector payload is {} bytes, shorter than its {FLAT_BINARY_HEADER_SIZE}-byte header",
460                    handle.len()
461                ),
462            ));
463        }
464
465        // Read header: magic(4) + dim(4) + num_vectors(4) + quant_type(1) + pad(3) = 16 bytes
466        let header = handle.read_bytes_range(0..header_len).await?;
467        if header.len() != FLAT_BINARY_HEADER_SIZE {
468            return Err(io::Error::new(
469                io::ErrorKind::UnexpectedEof,
470                format!(
471                    "flat vector header read returned {} bytes, expected {FLAT_BINARY_HEADER_SIZE}",
472                    header.len()
473                ),
474            ));
475        }
476        let hdr = header.as_slice();
477
478        let magic = u32::from_le_bytes([hdr[0], hdr[1], hdr[2], hdr[3]]);
479        if magic != FLAT_BINARY_MAGIC {
480            return Err(io::Error::new(
481                io::ErrorKind::InvalidData,
482                "Invalid FlatVectorData binary magic",
483            ));
484        }
485
486        let dim = u32::from_le_bytes([hdr[4], hdr[5], hdr[6], hdr[7]]) as usize;
487        let num_vectors = u32::from_le_bytes([hdr[8], hdr[9], hdr[10], hdr[11]]) as usize;
488        let quantization = DenseVectorQuantization::from_tag(hdr[12]).ok_or_else(|| {
489            io::Error::new(
490                io::ErrorKind::InvalidData,
491                format!("Unknown quantization tag: {}", hdr[12]),
492            )
493        })?;
494        if hdr[13..] != [0, 0, 0] {
495            return Err(io::Error::new(
496                io::ErrorKind::InvalidData,
497                "flat vector header has non-zero reserved bytes",
498            ));
499        }
500
501        // Read doc_ids section as zero-copy OwnedBytes (6 bytes per vector)
502        let vbs =
503            FlatVectorData::validate_shape(dim, num_vectors, quantization).map_err(|error| {
504                io::Error::new(
505                    io::ErrorKind::InvalidData,
506                    format!("invalid flat vector shape: {error}"),
507                )
508            })?;
509        let vectors_byte_len_usize = num_vectors.checked_mul(vbs).ok_or_else(|| {
510            io::Error::new(
511                io::ErrorKind::InvalidData,
512                "flat vector payload size overflows usize",
513            )
514        })?;
515        let doc_ids_byte_len_usize =
516            num_vectors.checked_mul(DOC_ID_ENTRY_SIZE).ok_or_else(|| {
517                io::Error::new(
518                    io::ErrorKind::InvalidData,
519                    "flat vector doc-map size overflows usize",
520                )
521            })?;
522        let expected_len_usize = FLAT_BINARY_HEADER_SIZE
523            .checked_add(vectors_byte_len_usize)
524            .and_then(|size| size.checked_add(doc_ids_byte_len_usize))
525            .ok_or_else(|| {
526                io::Error::new(
527                    io::ErrorKind::InvalidData,
528                    "flat vector serialized size overflows usize",
529                )
530            })?;
531        let expected_len = u64::try_from(expected_len_usize).map_err(|_| {
532            io::Error::new(
533                io::ErrorKind::InvalidData,
534                "flat vector serialized size does not fit in u64",
535            )
536        })?;
537        if handle.len() != expected_len {
538            return Err(io::Error::new(
539                io::ErrorKind::InvalidData,
540                format!(
541                    "flat vector payload has {} bytes, expected exactly {expected_len}",
542                    handle.len()
543                ),
544            ));
545        }
546
547        let vectors_byte_len = u64::try_from(vectors_byte_len_usize).map_err(|_| {
548            io::Error::new(
549                io::ErrorKind::InvalidData,
550                "flat vector payload size does not fit in u64",
551            )
552        })?;
553        let doc_ids_byte_len = u64::try_from(doc_ids_byte_len_usize).map_err(|_| {
554            io::Error::new(
555                io::ErrorKind::InvalidData,
556                "flat vector doc-map size does not fit in u64",
557            )
558        })?;
559        let doc_ids_start = header_len.checked_add(vectors_byte_len).ok_or_else(|| {
560            io::Error::new(
561                io::ErrorKind::InvalidData,
562                "flat vector doc-map offset overflows u64",
563            )
564        })?;
565        let doc_ids_end = doc_ids_start.checked_add(doc_ids_byte_len).ok_or_else(|| {
566            io::Error::new(
567                io::ErrorKind::InvalidData,
568                "flat vector doc-map range overflows u64",
569            )
570        })?;
571
572        let doc_ids_bytes = if load_doc_map {
573            let bytes = handle.read_bytes_range(doc_ids_start..doc_ids_end).await?;
574            if bytes.len() != doc_ids_byte_len_usize {
575                return Err(io::Error::new(
576                    io::ErrorKind::UnexpectedEof,
577                    format!(
578                        "flat vector doc-map read returned {} bytes, expected {doc_ids_byte_len_usize}",
579                        bytes.len()
580                    ),
581                ));
582            }
583            bytes
584        } else {
585            OwnedBytes::empty()
586        };
587
588        let mut previous = None;
589        let mut num_docs_with_vectors = 0usize;
590        for entry in doc_ids_bytes.as_slice().chunks_exact(DOC_ID_ENTRY_SIZE) {
591            let doc_id = u32::from_le_bytes([entry[0], entry[1], entry[2], entry[3]]);
592            let ordinal = u16::from_le_bytes([entry[4], entry[5]]);
593            let current = (doc_id, ordinal);
594            if let Some(previous) = previous
595                && previous >= current
596            {
597                return Err(io::Error::new(
598                    io::ErrorKind::InvalidData,
599                    format!(
600                        "flat vector doc map must be strictly sorted by (doc_id, ordinal), found {previous:?} before {current:?}"
601                    ),
602                ));
603            }
604            if let Some(limit) = total_docs
605                && doc_id >= limit
606            {
607                return Err(io::Error::new(
608                    io::ErrorKind::InvalidData,
609                    format!(
610                        "flat vector doc map references document {doc_id}, but segment contains only {} documents",
611                        limit
612                    ),
613                ));
614            }
615            if previous.is_none_or(|(previous_doc_id, _)| previous_doc_id != doc_id) {
616                num_docs_with_vectors = num_docs_with_vectors.checked_add(1).ok_or_else(|| {
617                    io::Error::new(
618                        io::ErrorKind::InvalidData,
619                        "flat vector distinct-document count overflows usize",
620                    )
621                })?;
622            }
623            previous = Some(current);
624        }
625
626        Ok(Self {
627            dim,
628            num_vectors,
629            num_docs_with_vectors,
630            quantization,
631            doc_ids_bytes,
632            handle,
633            vectors_offset: header_len,
634            vbs,
635            vectors_byte_len,
636        })
637    }
638
639    fn checked_vector_range(
640        &self,
641        start_idx: usize,
642        count: usize,
643    ) -> io::Result<(std::ops::Range<u64>, usize)> {
644        let end_idx = start_idx.checked_add(count).ok_or_else(|| {
645            io::Error::new(
646                io::ErrorKind::InvalidInput,
647                "flat vector index range overflows usize",
648            )
649        })?;
650        if end_idx > self.num_vectors {
651            return Err(io::Error::new(
652                io::ErrorKind::InvalidInput,
653                format!(
654                    "flat vector range {start_idx}..{end_idx} exceeds {} vectors",
655                    self.num_vectors
656                ),
657            ));
658        }
659
660        let relative_offset = start_idx.checked_mul(self.vbs).ok_or_else(|| {
661            io::Error::new(
662                io::ErrorKind::InvalidData,
663                "flat vector byte offset overflows usize",
664            )
665        })?;
666        let byte_len = count.checked_mul(self.vbs).ok_or_else(|| {
667            io::Error::new(
668                io::ErrorKind::InvalidInput,
669                "flat vector byte length overflows usize",
670            )
671        })?;
672        let relative_offset = u64::try_from(relative_offset).map_err(|_| {
673            io::Error::new(
674                io::ErrorKind::InvalidData,
675                "flat vector byte offset does not fit in u64",
676            )
677        })?;
678        let byte_len_u64 = u64::try_from(byte_len).map_err(|_| {
679            io::Error::new(
680                io::ErrorKind::InvalidInput,
681                "flat vector byte length does not fit in u64",
682            )
683        })?;
684        let start = self
685            .vectors_offset
686            .checked_add(relative_offset)
687            .ok_or_else(|| {
688                io::Error::new(
689                    io::ErrorKind::InvalidData,
690                    "flat vector byte offset overflows u64",
691                )
692            })?;
693        let end = start.checked_add(byte_len_u64).ok_or_else(|| {
694            io::Error::new(
695                io::ErrorKind::InvalidData,
696                "flat vector byte range overflows u64",
697            )
698        })?;
699        let vectors_end = self
700            .vectors_offset
701            .checked_add(self.vectors_byte_len)
702            .ok_or_else(|| {
703                io::Error::new(
704                    io::ErrorKind::InvalidData,
705                    "flat vector payload boundary overflows u64",
706                )
707            })?;
708        if end > vectors_end || end > self.handle.len() {
709            return Err(io::Error::new(
710                io::ErrorKind::InvalidData,
711                format!(
712                    "flat vector byte range {start}..{end} exceeds payload boundary {vectors_end}"
713                ),
714            ));
715        }
716        Ok((start..end, byte_len))
717    }
718
719    /// Pin the doc-id map (priority 3: every rerank / top-k resolution
720    /// binary-searches it).
721    #[cfg(feature = "native")]
722    pub(crate) fn pin_doc_ids(
723        &mut self,
724        mode: crate::segment::pin::PinMode,
725        remaining: &mut u64,
726        report: &mut crate::segment::pin::PinReport,
727    ) {
728        crate::segment::pin::pin_section(
729            &mut self.doc_ids_bytes,
730            "flat doc_ids",
731            mode,
732            remaining,
733            report,
734        );
735    }
736
737    /// Advise the kernel that vector data will be accessed at random offsets.
738    ///
739    /// Disables kernel readahead for the raw vector region. Rerank reads
740    /// scattered ~vbs-sized records; default readahead pulls in 128KB per
741    /// fault, evicting useful pages in memory-bound environments.
742    /// No-op for non-mmap (RAM, HTTP) backing.
743    #[cfg(feature = "native")]
744    pub fn advise_random_access(&self) {
745        let Some(vectors_end) = self.vectors_offset.checked_add(self.vectors_byte_len) else {
746            return;
747        };
748        self.handle
749            .madvise_range(self.vectors_offset..vectors_end, libc::MADV_RANDOM);
750    }
751
752    /// Prefetch the pages backing a sorted set of vector indexes (`MADV_WILLNEED`).
753    ///
754    /// Coalesces adjacent candidates into ranges so the kernel can overlap
755    /// the page-ins instead of taking one synchronous major fault per vector
756    /// during the rerank read loop. Indexes must be yielded in ascending order.
757    /// No-op for non-mmap backing.
758    #[cfg(feature = "native")]
759    pub fn prefetch_vectors(&self, sorted_flat_indexes: impl IntoIterator<Item = usize>) {
760        /// Gap (in bytes) below which two candidate ranges are merged into one advice call.
761        const COALESCE_GAP: u64 = 64 * 1024;
762        let mut ranges = sorted_flat_indexes.into_iter().filter_map(|idx| {
763            self.checked_vector_range(idx, 1)
764                .ok()
765                .map(|(range, _)| range)
766        });
767        let Some(first) = ranges.next() else {
768            return;
769        };
770        let mut run_start = first.start;
771        let mut run_end = first.end;
772        for range in ranges {
773            if range.start <= run_end.saturating_add(COALESCE_GAP) {
774                run_end = run_end.max(range.end);
775            } else {
776                self.handle
777                    .madvise_range(run_start..run_end, libc::MADV_WILLNEED);
778                run_start = range.start;
779                run_end = range.end;
780            }
781        }
782        self.handle
783            .madvise_range(run_start..run_end, libc::MADV_WILLNEED);
784    }
785
786    /// Read a single vector by index, dequantized to f32.
787    ///
788    /// `out` must have length >= `self.dim`. Returns `Ok(())` on success.
789    /// Used for ANN training and doc() hydration where f32 is needed.
790    pub async fn read_vector_into(&self, idx: usize, out: &mut [f32]) -> io::Result<()> {
791        if out.len() < self.dim {
792            return Err(io::Error::new(
793                io::ErrorKind::InvalidInput,
794                format!(
795                    "flat vector output is too short: need {} floats, got {}",
796                    self.dim,
797                    out.len()
798                ),
799            ));
800        }
801        let bytes = self.read_vectors_batch(idx, 1).await?;
802        dequantize_raw(bytes.as_slice(), self.quantization, self.dim, out)
803    }
804
805    /// Read a single vector by index, dequantized to f32 (allocates a new Vec<f32>).
806    pub async fn get_vector(&self, idx: usize) -> io::Result<Vec<f32>> {
807        let mut vector = vec![0f32; self.dim];
808        self.read_vector_into(idx, &mut vector).await?;
809        Ok(vector)
810    }
811
812    /// Read a single vector's raw bytes (no dequantization) into a caller-provided buffer.
813    ///
814    /// `out` must have length >= `self.vector_byte_size()`.
815    /// Used for native-precision reranking where raw quantized bytes are scored directly.
816    pub async fn read_vector_raw_into(&self, idx: usize, out: &mut [u8]) -> io::Result<()> {
817        self.read_vector_prefix_raw_into(idx, self.vector_byte_size(), out)
818            .await
819    }
820
821    /// Read a prefix of one vector's raw bytes into a caller-provided buffer.
822    ///
823    /// This is used by Matryoshka scoring to avoid reading the unused tail of
824    /// a vector. Unlike the old full-vector boundary, all caller-controlled
825    /// sizes and offset arithmetic are checked in release builds.
826    pub async fn read_vector_prefix_raw_into(
827        &self,
828        idx: usize,
829        prefix_byte_len: usize,
830        out: &mut [u8],
831    ) -> io::Result<()> {
832        let vbs = self.vector_byte_size();
833        if prefix_byte_len > vbs {
834            return Err(io::Error::new(
835                io::ErrorKind::InvalidInput,
836                format!(
837                    "vector prefix is {prefix_byte_len} bytes, but a vector has only {vbs} bytes"
838                ),
839            ));
840        }
841        if out.len() < prefix_byte_len {
842            return Err(io::Error::new(
843                io::ErrorKind::InvalidInput,
844                format!(
845                    "vector prefix output is too short: need {prefix_byte_len} bytes, got {}",
846                    out.len()
847                ),
848            ));
849        }
850        let (full_range, _) = self.checked_vector_range(idx, 1)?;
851        if prefix_byte_len == 0 {
852            return Ok(());
853        }
854        let prefix_byte_len_u64 = u64::try_from(prefix_byte_len).map_err(|_| {
855            io::Error::new(
856                io::ErrorKind::InvalidInput,
857                "vector prefix length does not fit in u64",
858            )
859        })?;
860        let byte_end = full_range
861            .start
862            .checked_add(prefix_byte_len_u64)
863            .ok_or_else(|| {
864                io::Error::new(
865                    io::ErrorKind::InvalidData,
866                    "vector byte range overflows u64",
867                )
868            })?;
869        let bytes = self
870            .handle
871            .read_bytes_range(full_range.start..byte_end)
872            .await?;
873        if bytes.len() != prefix_byte_len {
874            return Err(io::Error::new(
875                io::ErrorKind::UnexpectedEof,
876                format!(
877                    "vector prefix read returned {} bytes, expected {prefix_byte_len}",
878                    bytes.len()
879                ),
880            ));
881        }
882        out[..prefix_byte_len].copy_from_slice(bytes.as_slice());
883        Ok(())
884    }
885
886    /// Read a contiguous batch of raw quantized bytes by index range.
887    ///
888    /// Returns raw bytes for vectors `[start_idx..start_idx+count)`.
889    /// Bytes are in native quantized format — pass to `batch_cosine_scores_f16/u8`
890    /// or `batch_cosine_scores` (for f32) for scoring.
891    pub async fn read_vectors_batch(
892        &self,
893        start_idx: usize,
894        count: usize,
895    ) -> io::Result<OwnedBytes> {
896        let (range, expected_len) = self.checked_vector_range(start_idx, count)?;
897        let bytes = self.handle.read_bytes_range(range).await?;
898        if bytes.len() != expected_len {
899            return Err(io::Error::new(
900                io::ErrorKind::UnexpectedEof,
901                format!(
902                    "flat vector batch read returned {} bytes, expected {expected_len}",
903                    bytes.len()
904                ),
905            ));
906        }
907        Ok(bytes)
908    }
909
910    /// Synchronous read of a single vector's raw bytes.
911    #[cfg(feature = "sync")]
912    pub fn read_vector_raw_into_sync(&self, idx: usize, out: &mut [u8]) -> io::Result<()> {
913        let vbs = self.vector_byte_size();
914        if out.len() < vbs {
915            return Err(io::Error::new(
916                io::ErrorKind::InvalidInput,
917                format!(
918                    "flat vector output is too short: need {vbs} bytes, got {}",
919                    out.len()
920                ),
921            ));
922        }
923        let bytes = self.read_vectors_batch_sync(idx, 1)?;
924        out[..vbs].copy_from_slice(bytes.as_slice());
925        Ok(())
926    }
927
928    /// Synchronous batch read of raw quantized bytes.
929    #[cfg(feature = "sync")]
930    pub fn read_vectors_batch_sync(
931        &self,
932        start_idx: usize,
933        count: usize,
934    ) -> io::Result<OwnedBytes> {
935        let (range, expected_len) = self.checked_vector_range(start_idx, count)?;
936        let bytes = self.handle.read_bytes_range_sync(range)?;
937        if bytes.len() != expected_len {
938            return Err(io::Error::new(
939                io::ErrorKind::UnexpectedEof,
940                format!(
941                    "flat vector batch read returned {} bytes, expected {expected_len}",
942                    bytes.len()
943                ),
944            ));
945        }
946        Ok(bytes)
947    }
948
949    /// Find flat index range for a given doc_id (non-allocating).
950    ///
951    /// Returns `(start_index, count)` — the flat vector index range for this doc_id.
952    /// Use `get_doc_id(start + i)` for `i in 0..count` to read individual entries.
953    /// More efficient than `flat_indexes_for_doc` as it avoids Vec allocation.
954    pub fn flat_indexes_for_doc_range(&self, doc_id: u32) -> (usize, usize) {
955        let n = self.num_vectors;
956        let start = {
957            let mut lo = 0usize;
958            let mut hi = n;
959            while lo < hi {
960                let mid = lo + (hi - lo) / 2;
961                if self.doc_id_at(mid) < doc_id {
962                    lo = mid + 1;
963                } else {
964                    hi = mid;
965                }
966            }
967            lo
968        };
969        let mut count = 0;
970        let mut i = start;
971        while i < n && self.doc_id_at(i) == doc_id {
972            count += 1;
973            i += 1;
974        }
975        (start, count)
976    }
977
978    /// Find flat indexes for a given doc_id via binary search on sorted doc_ids.
979    ///
980    /// doc_ids are sorted by (doc_id, ordinal) — segment builder adds docs
981    /// sequentially. Binary search runs directly on zero-copy mmap bytes.
982    ///
983    /// Returns `(start_index, entries)` where start_index is the flat vector index.
984    pub fn flat_indexes_for_doc(&self, doc_id: u32) -> (usize, Vec<(u32, u16)>) {
985        let n = self.num_vectors;
986        // Binary search: find first entry where doc_id >= target
987        let start = {
988            let mut lo = 0usize;
989            let mut hi = n;
990            while lo < hi {
991                let mid = lo + (hi - lo) / 2;
992                if self.doc_id_at(mid) < doc_id {
993                    lo = mid + 1;
994                } else {
995                    hi = mid;
996                }
997            }
998            lo
999        };
1000        // Collect entries with matching doc_id
1001        let mut entries = Vec::new();
1002        let mut i = start;
1003        while i < n {
1004            let (did, ord) = self.get_doc_id(i);
1005            if did != doc_id {
1006                break;
1007            }
1008            entries.push((did, ord));
1009            i += 1;
1010        }
1011        (start, entries)
1012    }
1013
1014    /// Read doc_id at index from raw bytes (no ordinal).
1015    #[inline]
1016    fn doc_id_at(&self, idx: usize) -> u32 {
1017        assert_eq!(
1018            self.doc_ids_bytes.len(),
1019            self.num_vectors * DOC_ID_ENTRY_SIZE,
1020            "document IDs are unavailable on a training-only flat-vector reader",
1021        );
1022        let off = idx * DOC_ID_ENTRY_SIZE;
1023        let d = &self.doc_ids_bytes[off..];
1024        u32::from_le_bytes([d[0], d[1], d[2], d[3]])
1025    }
1026
1027    /// Get doc_id and ordinal at index (parsed from zero-copy mmap bytes).
1028    #[inline]
1029    pub fn get_doc_id(&self, idx: usize) -> (u32, u16) {
1030        assert_eq!(
1031            self.doc_ids_bytes.len(),
1032            self.num_vectors * DOC_ID_ENTRY_SIZE,
1033            "document IDs are unavailable on a training-only flat-vector reader",
1034        );
1035        let off = idx * DOC_ID_ENTRY_SIZE;
1036        let d = &self.doc_ids_bytes[off..];
1037        let doc_id = u32::from_le_bytes([d[0], d[1], d[2], d[3]]);
1038        let ordinal = u16::from_le_bytes([d[4], d[5]]);
1039        (doc_id, ordinal)
1040    }
1041
1042    /// Bytes per vector in storage (cached).
1043    #[inline]
1044    pub fn vector_byte_size(&self) -> usize {
1045        self.vbs
1046    }
1047
1048    /// Number of distinct documents that have at least one vector in this field.
1049    #[inline]
1050    pub fn num_docs_with_vectors(&self) -> usize {
1051        self.num_docs_with_vectors
1052    }
1053
1054    /// Total byte length of raw vector data (for chunked merger streaming).
1055    pub fn vector_bytes_len(&self) -> u64 {
1056        self.vectors_byte_len
1057    }
1058
1059    /// Byte offset where vector data starts (for direct handle access in merger).
1060    pub fn vectors_byte_offset(&self) -> u64 {
1061        self.vectors_offset
1062    }
1063
1064    /// Access the underlying file handle (for chunked byte-range reads in merger).
1065    pub fn handle(&self) -> &FileHandle {
1066        &self.handle
1067    }
1068
1069    /// Estimated memory usage — doc_ids are mmap-backed (only Arc overhead).
1070    pub fn estimated_memory_bytes(&self) -> usize {
1071        size_of::<Self>() + size_of::<OwnedBytes>()
1072    }
1073}
1074
1075#[cfg(test)]
1076mod tests {
1077    use super::*;
1078
1079    #[test]
1080    fn dequantize_raw_accepts_valid_storage_formats() {
1081        let f32_values = [1.25f32, -2.5];
1082        let f32_bytes = unsafe {
1083            // Safety: viewing an initialized f32 array as bytes is always valid.
1084            std::slice::from_raw_parts(
1085                f32_values.as_ptr().cast::<u8>(),
1086                std::mem::size_of_val(&f32_values),
1087            )
1088        };
1089        let mut out = [0.0; 2];
1090        dequantize_raw(f32_bytes, DenseVectorQuantization::F32, 2, &mut out).unwrap();
1091        assert_eq!(out, f32_values);
1092
1093        let f16_values = [0x3c00u16, 0xc000u16]; // 1.0, -2.0
1094        let f16_bytes = unsafe {
1095            // Safety: viewing an initialized u16 array as bytes is always valid.
1096            std::slice::from_raw_parts(
1097                f16_values.as_ptr().cast::<u8>(),
1098                std::mem::size_of_val(&f16_values),
1099            )
1100        };
1101        dequantize_raw(f16_bytes, DenseVectorQuantization::F16, 2, &mut out).unwrap();
1102        assert_eq!(out, [1.0, -2.0]);
1103
1104        dequantize_raw(&[0, u8::MAX], DenseVectorQuantization::UInt8, 2, &mut out).unwrap();
1105        assert_eq!(out, [u8_to_f32(0), u8_to_f32(u8::MAX)]);
1106    }
1107
1108    #[test]
1109    fn dequantize_raw_rejects_invalid_lengths_and_binary_storage() {
1110        let mut out = [0.0; 2];
1111        assert_eq!(
1112            dequantize_raw(&[0; 7], DenseVectorQuantization::F32, 2, &mut out)
1113                .unwrap_err()
1114                .kind(),
1115            io::ErrorKind::InvalidData
1116        );
1117        assert_eq!(
1118            dequantize_raw(&[0; 8], DenseVectorQuantization::F32, 2, &mut out[..1])
1119                .unwrap_err()
1120                .kind(),
1121            io::ErrorKind::InvalidInput
1122        );
1123        assert_eq!(
1124            dequantize_raw(&[], DenseVectorQuantization::Binary, 0, &mut [])
1125                .unwrap_err()
1126                .kind(),
1127            io::ErrorKind::InvalidInput
1128        );
1129    }
1130
1131    #[test]
1132    fn dequantize_raw_rejects_misaligned_typed_storage() {
1133        let storage = [0u8; 9];
1134        let offset = if (storage.as_ptr() as usize).is_multiple_of(4) {
1135            1
1136        } else {
1137            0
1138        };
1139        let raw = &storage[offset..offset + 8];
1140        assert!(!(raw.as_ptr() as usize).is_multiple_of(4));
1141
1142        let mut out = [0.0; 2];
1143        assert_eq!(
1144            dequantize_raw(raw, DenseVectorQuantization::F32, 2, &mut out)
1145                .unwrap_err()
1146                .kind(),
1147            io::ErrorKind::InvalidData
1148        );
1149    }
1150
1151    #[test]
1152    fn flat_vector_writers_reject_inconsistent_shapes_and_doc_maps() {
1153        let mut encoded = Vec::new();
1154        assert!(
1155            FlatVectorData::serialize_binary_from_flat_streaming(
1156                0,
1157                &[],
1158                &[],
1159                DenseVectorQuantization::F32,
1160                &mut encoded,
1161            )
1162            .is_err()
1163        );
1164        assert!(encoded.is_empty());
1165
1166        assert!(
1167            FlatVectorData::serialize_binary_from_flat_streaming(
1168                2,
1169                &[1.0],
1170                &[(0, 0)],
1171                DenseVectorQuantization::F32,
1172                &mut encoded,
1173            )
1174            .is_err()
1175        );
1176        assert!(encoded.is_empty());
1177
1178        assert!(
1179            FlatVectorData::serialize_binary_from_flat_streaming(
1180                1,
1181                &[1.0],
1182                &[(0, 0)],
1183                DenseVectorQuantization::Binary,
1184                &mut encoded,
1185            )
1186            .is_err()
1187        );
1188        assert!(encoded.is_empty());
1189
1190        assert!(
1191            FlatVectorData::serialize_binary_from_flat_streaming(
1192                1,
1193                &[1.0, 2.0],
1194                &[(1, 0), (0, 0)],
1195                DenseVectorQuantization::F32,
1196                &mut encoded,
1197            )
1198            .is_err()
1199        );
1200        assert!(encoded.is_empty());
1201
1202        assert!(
1203            FlatVectorData::serialize_binary_from_flat_streaming(
1204                1,
1205                &[1.0, 2.0],
1206                &[(0, 0), (0, 0)],
1207                DenseVectorQuantization::F32,
1208                &mut encoded,
1209            )
1210            .is_err()
1211        );
1212        assert!(encoded.is_empty());
1213
1214        assert!(
1215            FlatVectorData::serialize_binary_from_bits_streaming(7, &[0], &[(0, 0)], &mut encoded,)
1216                .is_err()
1217        );
1218        assert!(encoded.is_empty());
1219
1220        assert!(
1221            FlatVectorData::serialize_binary_from_bits_streaming(8, &[], &[(0, 0)], &mut encoded,)
1222                .is_err()
1223        );
1224        assert!(encoded.is_empty());
1225
1226        let vectors = [1.0f32, 2.0, 3.0, 4.0];
1227        let doc_ids = [(0, 0), (1, 0)];
1228        FlatVectorData::serialize_binary_from_flat_streaming(
1229            2,
1230            &vectors,
1231            &doc_ids,
1232            DenseVectorQuantization::F32,
1233            &mut encoded,
1234        )
1235        .unwrap();
1236        assert_eq!(
1237            encoded.len(),
1238            FlatVectorData::serialized_binary_size(2, 2, DenseVectorQuantization::F32).unwrap()
1239        );
1240    }
1241
1242    fn encoded_two_vector_payload() -> Vec<u8> {
1243        let mut encoded = Vec::new();
1244        FlatVectorData::serialize_binary_from_flat_streaming(
1245            2,
1246            &[1.0, 2.0, 3.0, 4.0],
1247            &[(0, 0), (1, 0)],
1248            DenseVectorQuantization::F32,
1249            &mut encoded,
1250        )
1251        .unwrap();
1252        encoded
1253    }
1254
1255    #[tokio::test]
1256    async fn flat_vector_open_rejects_corrupt_layout_and_doc_map() {
1257        let valid = encoded_two_vector_payload();
1258
1259        let mut multi_value = Vec::new();
1260        FlatVectorData::serialize_binary_from_flat_streaming(
1261            1,
1262            &[1.0, 2.0, 3.0],
1263            &[(0, 0), (0, 1), (2, 0)],
1264            DenseVectorQuantization::F32,
1265            &mut multi_value,
1266        )
1267        .unwrap();
1268        let multi_value = LazyFlatVectorData::open_with_doc_limit(
1269            FileHandle::from_bytes(OwnedBytes::new(multi_value)),
1270            Some(3),
1271        )
1272        .await
1273        .unwrap();
1274        assert_eq!(multi_value.num_docs_with_vectors(), 2);
1275
1276        let mut trailing = valid.clone();
1277        trailing.push(0);
1278        assert!(
1279            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(trailing)))
1280                .await
1281                .is_err()
1282        );
1283
1284        let mut truncated = valid.clone();
1285        truncated.pop();
1286        assert!(
1287            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(truncated)))
1288                .await
1289                .is_err()
1290        );
1291
1292        let mut reserved = valid.clone();
1293        reserved[13] = 1;
1294        assert!(
1295            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(reserved)))
1296                .await
1297                .is_err()
1298        );
1299
1300        let doc_map_start = FLAT_BINARY_HEADER_SIZE + 2 * 2 * size_of::<f32>();
1301        let mut unsorted = valid.clone();
1302        let (first, second) = unsorted[doc_map_start..doc_map_start + 2 * DOC_ID_ENTRY_SIZE]
1303            .split_at_mut(DOC_ID_ENTRY_SIZE);
1304        first.swap_with_slice(second);
1305        assert!(
1306            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(unsorted)))
1307                .await
1308                .is_err()
1309        );
1310
1311        let mut duplicate = valid.clone();
1312        duplicate.copy_within(
1313            doc_map_start..doc_map_start + DOC_ID_ENTRY_SIZE,
1314            doc_map_start + DOC_ID_ENTRY_SIZE,
1315        );
1316        assert!(
1317            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(duplicate)))
1318                .await
1319                .is_err()
1320        );
1321
1322        assert!(
1323            LazyFlatVectorData::open_with_doc_limit(
1324                FileHandle::from_bytes(OwnedBytes::new(valid)),
1325                Some(1),
1326            )
1327            .await
1328            .is_err()
1329        );
1330
1331        let mut invalid_binary = Vec::new();
1332        FlatVectorData::serialize_binary_from_bits_streaming(
1333            8,
1334            &[0],
1335            &[(0, 0)],
1336            &mut invalid_binary,
1337        )
1338        .unwrap();
1339        invalid_binary[4..8].copy_from_slice(&7u32.to_le_bytes());
1340        assert!(
1341            LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(invalid_binary)))
1342                .await
1343                .is_err()
1344        );
1345    }
1346
1347    #[tokio::test]
1348    async fn flat_vector_batch_and_dequantized_reads_are_checked() {
1349        let flat = LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(
1350            encoded_two_vector_payload(),
1351        )))
1352        .await
1353        .unwrap();
1354
1355        assert_eq!(flat.read_vectors_batch(0, 2).await.unwrap().len(), 16);
1356        assert_eq!(flat.read_vectors_batch(2, 0).await.unwrap().len(), 0);
1357        assert!(flat.read_vectors_batch(1, 2).await.is_err());
1358        assert!(flat.read_vectors_batch(usize::MAX, 1).await.is_err());
1359        assert!(flat.read_vectors_batch(0, usize::MAX).await.is_err());
1360
1361        let mut values = [0.0; 2];
1362        flat.read_vector_into(1, &mut values).await.unwrap();
1363        assert_eq!(values, [3.0, 4.0]);
1364        assert!(flat.read_vector_into(2, &mut values).await.is_err());
1365        assert!(flat.read_vector_into(0, &mut values[..1]).await.is_err());
1366
1367        #[cfg(feature = "sync")]
1368        {
1369            assert_eq!(flat.read_vectors_batch_sync(0, 2).unwrap().len(), 16);
1370            assert!(flat.read_vectors_batch_sync(1, 2).is_err());
1371            assert!(flat.read_vectors_batch_sync(usize::MAX, 1).is_err());
1372            let mut too_short = [0; 7];
1373            assert!(flat.read_vector_raw_into_sync(0, &mut too_short).is_err());
1374        }
1375    }
1376
1377    #[cfg(not(target_arch = "wasm32"))]
1378    #[tokio::test]
1379    async fn flat_vector_reads_reject_short_lazy_range_results() {
1380        let payload = std::sync::Arc::new(encoded_two_vector_payload());
1381        let payload_len = payload.len() as u64;
1382        let read_payload = std::sync::Arc::clone(&payload);
1383        let read_fn: crate::directories::RangeReadFn = std::sync::Arc::new(move |range| {
1384            let payload = std::sync::Arc::clone(&read_payload);
1385            Box::pin(async move {
1386                let start = usize::try_from(range.start).unwrap();
1387                let mut end = usize::try_from(range.end).unwrap();
1388                // Header and doc-map reads are exact, allowing open to finish.
1389                // Raw vector reads deliberately violate the range-read contract.
1390                if range.start == FLAT_BINARY_HEADER_SIZE as u64 {
1391                    end -= 1;
1392                }
1393                Ok(OwnedBytes::new(payload[start..end].to_vec()))
1394            })
1395        });
1396        let flat = LazyFlatVectorData::open(FileHandle::lazy(payload_len, read_fn))
1397            .await
1398            .unwrap();
1399
1400        let error = flat.read_vectors_batch(0, 1).await.unwrap_err();
1401        assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
1402        let mut raw = [0; 8];
1403        let error = flat.read_vector_raw_into(0, &mut raw).await.unwrap_err();
1404        assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
1405    }
1406
1407    #[tokio::test]
1408    async fn vector_prefix_reads_are_checked_and_do_not_fetch_the_tail() {
1409        let vectors = [1.0f32, 2.0, 3.0, 4.0];
1410        let doc_ids = [(0, 0), (1, 0)];
1411        let mut encoded = Vec::new();
1412        FlatVectorData::serialize_binary_from_flat_streaming(
1413            2,
1414            &vectors,
1415            &doc_ids,
1416            DenseVectorQuantization::F32,
1417            &mut encoded,
1418        )
1419        .unwrap();
1420        let flat = LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(encoded)))
1421            .await
1422            .unwrap();
1423
1424        let mut prefix = [0xa5; 8];
1425        flat.read_vector_prefix_raw_into(1, 4, &mut prefix)
1426            .await
1427            .unwrap();
1428        assert_eq!(&prefix[..4], &3.0f32.to_ne_bytes());
1429        assert_eq!(&prefix[4..], &[0xa5; 4]);
1430
1431        let mut full = [0; 8];
1432        flat.read_vector_raw_into(1, &mut full).await.unwrap();
1433        assert_eq!(&full[..4], &3.0f32.to_ne_bytes());
1434        assert_eq!(&full[4..], &4.0f32.to_ne_bytes());
1435
1436        assert!(
1437            flat.read_vector_prefix_raw_into(2, 4, &mut prefix)
1438                .await
1439                .is_err()
1440        );
1441        assert!(
1442            flat.read_vector_prefix_raw_into(0, 9, &mut prefix)
1443                .await
1444                .is_err()
1445        );
1446        assert!(
1447            flat.read_vector_prefix_raw_into(0, 4, &mut prefix[..3])
1448                .await
1449                .is_err()
1450        );
1451    }
1452}