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