1use 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#[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 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 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
102pub struct FlatVectorData;
116
117impl FlatVectorData {
118 fn validate_shape(
119 dim: usize,
120 num_vectors: usize,
121 quant: DenseVectorQuantization,
122 ) -> io::Result<usize> {
123 if dim == 0 {
124 return Err(io::Error::new(
125 io::ErrorKind::InvalidInput,
126 "flat vector dimension must be greater than zero",
127 ));
128 }
129 if quant == DenseVectorQuantization::Binary && !dim.is_multiple_of(8) {
130 return Err(io::Error::new(
131 io::ErrorKind::InvalidInput,
132 format!("binary flat vector dimension must be a multiple of 8, got {dim}"),
133 ));
134 }
135 u32::try_from(dim).map_err(|_| {
136 io::Error::new(
137 io::ErrorKind::InvalidInput,
138 format!("flat vector dimension {dim} exceeds u32::MAX"),
139 )
140 })?;
141 u32::try_from(num_vectors).map_err(|_| {
142 io::Error::new(
143 io::ErrorKind::InvalidInput,
144 format!("flat vector count {num_vectors} exceeds u32::MAX"),
145 )
146 })?;
147
148 match quant {
149 DenseVectorQuantization::Binary => dim.checked_add(7).map(|bits| bits / 8),
150 _ => dim.checked_mul(quant.element_size()),
151 }
152 .ok_or_else(|| {
153 io::Error::new(
154 io::ErrorKind::InvalidInput,
155 "flat vector byte size overflows usize",
156 )
157 })
158 }
159
160 fn validate_doc_ids(doc_ids: &[(u32, u16)]) -> io::Result<()> {
161 if let Some(pair) = doc_ids.windows(2).find(|pair| pair[0] >= pair[1]) {
162 return Err(io::Error::new(
163 io::ErrorKind::InvalidInput,
164 format!(
165 "flat vector doc map must be strictly sorted by (doc_id, ordinal), found {:?} before {:?}",
166 pair[0], pair[1]
167 ),
168 ));
169 }
170 Ok(())
171 }
172
173 pub(crate) fn validate_dense_input(
176 dim: usize,
177 flat_vectors: &[f32],
178 doc_ids: &[(u32, u16)],
179 quant: DenseVectorQuantization,
180 ) -> io::Result<usize> {
181 if quant == DenseVectorQuantization::Binary {
182 return Err(io::Error::new(
183 io::ErrorKind::InvalidInput,
184 "binary quantization must use serialize_binary_from_bits_streaming",
185 ));
186 }
187 let num_vectors = doc_ids.len();
188 let expected_floats = num_vectors.checked_mul(dim).ok_or_else(|| {
189 io::Error::new(
190 io::ErrorKind::InvalidInput,
191 "flat f32 vector count overflows usize",
192 )
193 })?;
194 if flat_vectors.len() != expected_floats {
195 return Err(io::Error::new(
196 io::ErrorKind::InvalidInput,
197 format!(
198 "flat vector input has {} floats, expected {num_vectors} x {dim} = {expected_floats}",
199 flat_vectors.len()
200 ),
201 ));
202 }
203 Self::validate_doc_ids(doc_ids)?;
204 Self::serialized_binary_size(dim, num_vectors, quant)
205 }
206
207 pub(crate) fn validate_binary_input(
210 dim_bits: usize,
211 packed_vectors: &[u8],
212 doc_ids: &[(u32, u16)],
213 ) -> io::Result<usize> {
214 let num_vectors = doc_ids.len();
215 let byte_len =
216 Self::validate_shape(dim_bits, num_vectors, DenseVectorQuantization::Binary)?;
217 let expected_bytes = num_vectors.checked_mul(byte_len).ok_or_else(|| {
218 io::Error::new(
219 io::ErrorKind::InvalidInput,
220 "packed binary vector size overflows usize",
221 )
222 })?;
223 if packed_vectors.len() != expected_bytes {
224 return Err(io::Error::new(
225 io::ErrorKind::InvalidInput,
226 format!(
227 "packed binary input has {} bytes, expected {num_vectors} x {byte_len} = {expected_bytes}",
228 packed_vectors.len()
229 ),
230 ));
231 }
232 Self::validate_doc_ids(doc_ids)?;
233 Self::serialized_binary_size(dim_bits, num_vectors, DenseVectorQuantization::Binary)
234 }
235
236 pub fn write_binary_header(
238 dim: usize,
239 num_vectors: usize,
240 quant: DenseVectorQuantization,
241 writer: &mut dyn std::io::Write,
242 ) -> std::io::Result<()> {
243 Self::validate_shape(dim, num_vectors, quant)?;
244 let dim = u32::try_from(dim).map_err(|_| {
245 io::Error::new(
246 io::ErrorKind::InvalidInput,
247 "flat vector dimension exceeds u32",
248 )
249 })?;
250 let num_vectors = u32::try_from(num_vectors).map_err(|_| {
251 io::Error::new(io::ErrorKind::InvalidInput, "flat vector count exceeds u32")
252 })?;
253 writer.write_all(&FLAT_BINARY_MAGIC.to_le_bytes())?;
254 writer.write_all(&dim.to_le_bytes())?;
255 writer.write_all(&num_vectors.to_le_bytes())?;
256 writer.write_all(&[quant.tag(), 0, 0, 0])?; Ok(())
258 }
259
260 pub fn serialized_binary_size(
262 dim: usize,
263 num_vectors: usize,
264 quant: DenseVectorQuantization,
265 ) -> io::Result<usize> {
266 let bytes_per_vector = Self::validate_shape(dim, num_vectors, quant)?;
267 let vector_bytes = num_vectors.checked_mul(bytes_per_vector).ok_or_else(|| {
268 io::Error::new(
269 io::ErrorKind::InvalidInput,
270 "flat vector payload size overflows usize",
271 )
272 })?;
273 let doc_id_bytes = num_vectors.checked_mul(DOC_ID_ENTRY_SIZE).ok_or_else(|| {
274 io::Error::new(
275 io::ErrorKind::InvalidInput,
276 "flat vector doc-map size overflows usize",
277 )
278 })?;
279 FLAT_BINARY_HEADER_SIZE
280 .checked_add(vector_bytes)
281 .and_then(|size| size.checked_add(doc_id_bytes))
282 .ok_or_else(|| {
283 io::Error::new(
284 io::ErrorKind::InvalidInput,
285 "flat vector serialized size overflows usize",
286 )
287 })
288 }
289
290 pub fn serialize_binary_from_flat_streaming(
295 dim: usize,
296 flat_vectors: &[f32],
297 doc_ids: &[(u32, u16)],
298 quant: DenseVectorQuantization,
299 writer: &mut dyn std::io::Write,
300 ) -> std::io::Result<()> {
301 Self::validate_dense_input(dim, flat_vectors, doc_ids, quant)?;
302 let num_vectors = doc_ids.len();
303 Self::write_binary_header(dim, num_vectors, quant, writer)?;
304
305 match quant {
306 DenseVectorQuantization::F32 => {
307 let bytes: &[u8] = unsafe {
308 std::slice::from_raw_parts(
309 flat_vectors.as_ptr() as *const u8,
310 std::mem::size_of_val(flat_vectors),
311 )
312 };
313 writer.write_all(bytes)?;
314 }
315 DenseVectorQuantization::F16 => {
316 let mut buf = vec![0u16; dim];
317 for v in flat_vectors.chunks_exact(dim) {
318 batch_f32_to_f16(v, &mut buf);
319 let bytes: &[u8] =
320 unsafe { std::slice::from_raw_parts(buf.as_ptr() as *const u8, dim * 2) };
321 writer.write_all(bytes)?;
322 }
323 }
324 DenseVectorQuantization::UInt8 => {
325 let mut buf = vec![0u8; dim];
326 for v in flat_vectors.chunks_exact(dim) {
327 batch_f32_to_u8(v, &mut buf);
328 writer.write_all(&buf)?;
329 }
330 }
331 DenseVectorQuantization::Binary => unreachable!("validated above"),
332 }
333
334 for &(doc_id, ordinal) in doc_ids {
335 writer.write_all(&doc_id.to_le_bytes())?;
336 writer.write_all(&ordinal.to_le_bytes())?;
337 }
338
339 Ok(())
340 }
341
342 pub fn serialize_binary_from_bits_streaming(
347 dim_bits: usize,
348 packed_vectors: &[u8],
349 doc_ids: &[(u32, u16)],
350 writer: &mut dyn std::io::Write,
351 ) -> std::io::Result<()> {
352 Self::validate_binary_input(dim_bits, packed_vectors, doc_ids)?;
353 let num_vectors = doc_ids.len();
354
355 Self::write_binary_header(
356 dim_bits,
357 num_vectors,
358 DenseVectorQuantization::Binary,
359 writer,
360 )?;
361 writer.write_all(packed_vectors)?;
362
363 for &(doc_id, ordinal) in doc_ids {
364 writer.write_all(&doc_id.to_le_bytes())?;
365 writer.write_all(&ordinal.to_le_bytes())?;
366 }
367
368 Ok(())
369 }
370
371 pub fn write_raw_vector_bytes(
375 raw_bytes: &[u8],
376 writer: &mut dyn std::io::Write,
377 ) -> std::io::Result<()> {
378 writer.write_all(raw_bytes)
379 }
380}
381
382#[derive(Debug, Clone)]
395pub struct LazyFlatVectorData {
396 pub dim: usize,
398 pub num_vectors: usize,
400 num_docs_with_vectors: usize,
402 pub quantization: DenseVectorQuantization,
404 doc_ids_bytes: OwnedBytes,
406 has_doc_map: bool,
410 handle: FileHandle,
412 vectors_offset: u64,
414 vbs: usize,
416 vectors_byte_len: u64,
418 locations: Option<super::vector_locations::VectorLocations>,
420}
421
422impl LazyFlatVectorData {
423 pub async fn open(handle: FileHandle) -> io::Result<Self> {
428 Self::open_with_doc_limit(handle, None).await
429 }
430
431 pub(crate) async fn open_with_doc_limit(
437 handle: FileHandle,
438 total_docs: Option<u32>,
439 ) -> io::Result<Self> {
440 Self::open_impl(handle, total_docs, true).await
441 }
442
443 pub(crate) async fn open_for_training(handle: FileHandle) -> io::Result<Self> {
449 Self::open_impl(handle, None, false).await
450 }
451
452 async fn open_impl(
453 handle: FileHandle,
454 total_docs: Option<u32>,
455 load_doc_map: bool,
456 ) -> io::Result<Self> {
457 let header_len = u64::try_from(FLAT_BINARY_HEADER_SIZE).map_err(|_| {
458 io::Error::new(
459 io::ErrorKind::InvalidData,
460 "flat vector header size does not fit in u64",
461 )
462 })?;
463 if handle.len() < header_len {
464 return Err(io::Error::new(
465 io::ErrorKind::UnexpectedEof,
466 format!(
467 "flat vector payload is {} bytes, shorter than its {FLAT_BINARY_HEADER_SIZE}-byte header",
468 handle.len()
469 ),
470 ));
471 }
472
473 let header = handle.read_bytes_range(0..header_len).await?;
475 if header.len() != FLAT_BINARY_HEADER_SIZE {
476 return Err(io::Error::new(
477 io::ErrorKind::UnexpectedEof,
478 format!(
479 "flat vector header read returned {} bytes, expected {FLAT_BINARY_HEADER_SIZE}",
480 header.len()
481 ),
482 ));
483 }
484 let hdr = header.as_slice();
485
486 let magic = u32::from_le_bytes([hdr[0], hdr[1], hdr[2], hdr[3]]);
487 if magic != FLAT_BINARY_MAGIC {
488 return Err(io::Error::new(
489 io::ErrorKind::InvalidData,
490 "Invalid FlatVectorData binary magic",
491 ));
492 }
493
494 let dim = u32::from_le_bytes([hdr[4], hdr[5], hdr[6], hdr[7]]) as usize;
495 let num_vectors = u32::from_le_bytes([hdr[8], hdr[9], hdr[10], hdr[11]]) as usize;
496 let quantization = DenseVectorQuantization::from_tag(hdr[12]).ok_or_else(|| {
497 io::Error::new(
498 io::ErrorKind::InvalidData,
499 format!("Unknown quantization tag: {}", hdr[12]),
500 )
501 })?;
502 if hdr[13..] != [0, 0, 0] {
503 return Err(io::Error::new(
504 io::ErrorKind::InvalidData,
505 "flat vector header has non-zero reserved bytes",
506 ));
507 }
508
509 let vbs =
511 FlatVectorData::validate_shape(dim, num_vectors, quantization).map_err(|error| {
512 io::Error::new(
513 io::ErrorKind::InvalidData,
514 format!("invalid flat vector shape: {error}"),
515 )
516 })?;
517 let vectors_byte_len_usize = num_vectors.checked_mul(vbs).ok_or_else(|| {
518 io::Error::new(
519 io::ErrorKind::InvalidData,
520 "flat vector payload size overflows usize",
521 )
522 })?;
523 let doc_ids_byte_len_usize =
524 num_vectors.checked_mul(DOC_ID_ENTRY_SIZE).ok_or_else(|| {
525 io::Error::new(
526 io::ErrorKind::InvalidData,
527 "flat vector doc-map size overflows usize",
528 )
529 })?;
530 let expected_len_usize = FLAT_BINARY_HEADER_SIZE
531 .checked_add(vectors_byte_len_usize)
532 .and_then(|size| size.checked_add(doc_ids_byte_len_usize))
533 .ok_or_else(|| {
534 io::Error::new(
535 io::ErrorKind::InvalidData,
536 "flat vector serialized size overflows usize",
537 )
538 })?;
539 let expected_len = u64::try_from(expected_len_usize).map_err(|_| {
540 io::Error::new(
541 io::ErrorKind::InvalidData,
542 "flat vector serialized size does not fit in u64",
543 )
544 })?;
545 if handle.len() != expected_len {
546 return Err(io::Error::new(
547 io::ErrorKind::InvalidData,
548 format!(
549 "flat vector payload has {} bytes, expected exactly {expected_len}",
550 handle.len()
551 ),
552 ));
553 }
554
555 let vectors_byte_len = u64::try_from(vectors_byte_len_usize).map_err(|_| {
556 io::Error::new(
557 io::ErrorKind::InvalidData,
558 "flat vector payload size does not fit in u64",
559 )
560 })?;
561 let doc_ids_byte_len = u64::try_from(doc_ids_byte_len_usize).map_err(|_| {
562 io::Error::new(
563 io::ErrorKind::InvalidData,
564 "flat vector doc-map size does not fit in u64",
565 )
566 })?;
567 let doc_ids_start = header_len.checked_add(vectors_byte_len).ok_or_else(|| {
568 io::Error::new(
569 io::ErrorKind::InvalidData,
570 "flat vector doc-map offset overflows u64",
571 )
572 })?;
573 let doc_ids_end = doc_ids_start.checked_add(doc_ids_byte_len).ok_or_else(|| {
574 io::Error::new(
575 io::ErrorKind::InvalidData,
576 "flat vector doc-map range overflows u64",
577 )
578 })?;
579
580 let doc_ids_bytes = if load_doc_map {
581 let bytes = handle.read_bytes_range(doc_ids_start..doc_ids_end).await?;
582 if bytes.len() != doc_ids_byte_len_usize {
583 return Err(io::Error::new(
584 io::ErrorKind::UnexpectedEof,
585 format!(
586 "flat vector doc-map read returned {} bytes, expected {doc_ids_byte_len_usize}",
587 bytes.len()
588 ),
589 ));
590 }
591 bytes
592 } else {
593 OwnedBytes::empty()
594 };
595
596 let num_docs_with_vectors =
597 Self::validate_doc_map(&doc_ids_bytes, DOC_ID_ENTRY_SIZE, total_docs, |doc, _| {
598 Ok(doc)
599 })?;
600
601 debug_assert!(!load_doc_map || doc_ids_bytes.len() == doc_ids_byte_len_usize);
602 Ok(Self {
603 dim,
604 num_vectors,
605 num_docs_with_vectors,
606 quantization,
607 doc_ids_bytes,
608 has_doc_map: load_doc_map,
609 handle,
610 vectors_offset: header_len,
611 vbs,
612 vectors_byte_len,
613 locations: None,
614 })
615 }
616
617 fn validate_doc_map(
618 bytes: &OwnedBytes,
619 stride: usize,
620 total_docs: Option<u32>,
621 mut resolve: impl FnMut(u32, u16) -> io::Result<u32>,
622 ) -> io::Result<usize> {
623 let mut previous = None;
624 let mut num_docs_with_vectors = 0usize;
625 for (index, entry) in bytes.chunks_exact(stride).enumerate() {
626 #[cfg(feature = "native")]
629 if index.is_multiple_of(64 * 1024) {
630 let start = index * stride;
631 let end = start.saturating_add(128 * 1024 * stride).min(bytes.len());
632 bytes.madvise_range(start..end, libc::MADV_WILLNEED);
633 }
634 #[cfg(not(feature = "native"))]
635 let _ = index;
636 let local_doc = u32::from_le_bytes([entry[0], entry[1], entry[2], entry[3]]);
637 let ordinal = u16::from_le_bytes([entry[4], entry[5]]);
638 let doc_id = resolve(local_doc, ordinal)?;
639 let current = (doc_id, ordinal);
640 if let Some(previous) = previous
641 && previous >= current
642 {
643 return Err(io::Error::new(
644 io::ErrorKind::InvalidData,
645 format!(
646 "flat vector doc map must be strictly sorted by (doc_id, ordinal), found {previous:?} before {current:?}"
647 ),
648 ));
649 }
650 if let Some(limit) = total_docs
651 && doc_id >= limit
652 {
653 return Err(io::Error::new(
654 io::ErrorKind::InvalidData,
655 format!(
656 "flat vector doc map references document {doc_id}, but segment contains only {} documents",
657 limit
658 ),
659 ));
660 }
661 if previous.is_none_or(|(previous_doc_id, _)| previous_doc_id != doc_id) {
662 num_docs_with_vectors = num_docs_with_vectors.checked_add(1).ok_or_else(|| {
663 io::Error::new(
664 io::ErrorKind::InvalidData,
665 "flat vector distinct-document count overflows usize",
666 )
667 })?;
668 }
669 previous = Some(current);
670 }
671
672 Ok(num_docs_with_vectors)
673 }
674
675 pub(crate) async fn open_indirect(
678 map: FileHandle,
679 ann: &crate::segment::ann_disk::AnnDiskIndex,
680 total_docs: u32,
681 ) -> io::Result<Self> {
682 use super::vector_locations::{ENTRY_SIZE, VectorLocations};
683 let bytes = map.read_bytes().await?;
684 if bytes.len() as u64 != map.len() {
685 return Err(io::Error::new(
686 io::ErrorKind::UnexpectedEof,
687 "truncated exact-vector lookup read",
688 ));
689 }
690 let opened = VectorLocations::open(bytes)?;
691 let handle = ann.exact_vectors_handle();
692 let dim = opened.dim;
693 let num_vectors = opened.count;
694 if opened.ann_len != handle.len() || dim != ann.header().dim {
695 return Err(io::Error::new(
696 io::ErrorKind::InvalidData,
697 "exact-vector lookup does not match its ANN payload",
698 ));
699 }
700 let vbs =
701 FlatVectorData::validate_shape(dim, num_vectors, DenseVectorQuantization::Binary)?;
702 let doc_ids_bytes = opened.rows;
703 let layout = opened.layout;
704 let num_docs_with_vectors = {
705 let spans = ann.exact_location_spans(layout.spans())?;
706 let mut addresses = layout.addresses(doc_ids_bytes.as_slice());
707 Self::validate_doc_map(
708 &doc_ids_bytes,
709 ENTRY_SIZE,
710 Some(total_docs),
711 |local_doc, ordinal| {
712 let (base, span, row) =
713 addresses.next().expect("validated lookup row count")?;
714 let doc = local_doc.checked_add(base).ok_or_else(|| {
715 io::Error::new(io::ErrorKind::InvalidData, "vector document base overflow")
716 })?;
717 spans[span].validate_label(row, doc, ordinal)?;
718 Ok(doc)
719 },
720 )?
721 };
722 let vectors_byte_len = (num_vectors as u64)
723 .checked_mul(vbs as u64)
724 .ok_or_else(|| {
725 io::Error::new(
726 io::ErrorKind::InvalidData,
727 "exact-vector byte count overflow",
728 )
729 })?;
730 Ok(Self {
731 dim,
732 num_vectors,
733 num_docs_with_vectors,
734 quantization: DenseVectorQuantization::Binary,
735 doc_ids_bytes,
736 has_doc_map: true,
737 handle,
738 vectors_offset: 0,
739 vbs,
740 vectors_byte_len,
741 locations: Some(layout),
742 })
743 }
744
745 pub fn is_ann_backed(&self) -> bool {
747 self.locations.is_some()
748 }
749
750 fn doc_map_stride(&self) -> usize {
751 if self.locations.is_some() {
752 super::vector_locations::ENTRY_SIZE
753 } else {
754 DOC_ID_ENTRY_SIZE
755 }
756 }
757
758 fn code_offset(&self, idx: usize) -> u64 {
759 let at = idx * super::vector_locations::ENTRY_SIZE + DOC_ID_ENTRY_SIZE;
760 let address = u64::from_le_bytes(self.doc_ids_bytes[at..at + 8].try_into().unwrap());
761 self.locations
762 .as_ref()
763 .expect("ANN-backed vector")
764 .code_offset(idx, address, self.vbs)
765 .expect("validated vector location")
766 }
767
768 #[cfg(feature = "native")]
769 pub(super) fn locations(&self) -> Option<&super::vector_locations::VectorLocations> {
770 self.locations.as_ref()
771 }
772
773 #[cfg(feature = "native")]
774 pub(super) fn location_rows(&self) -> &[u8] {
775 self.doc_ids_bytes.as_slice()
776 }
777
778 fn contiguous_rows(&self, start: usize, count: usize) -> usize {
780 if self.locations.is_none() || count < 2 {
781 return count;
782 }
783 let first = self.code_offset(start);
784 (1..count)
785 .find(|&i| self.code_offset(start + i) != first + i as u64 * self.vbs as u64)
786 .unwrap_or(count)
787 }
788
789 fn checked_vector_range(
790 &self,
791 start_idx: usize,
792 count: usize,
793 ) -> io::Result<(std::ops::Range<u64>, usize)> {
794 let end_idx = start_idx.checked_add(count).ok_or_else(|| {
795 io::Error::new(
796 io::ErrorKind::InvalidInput,
797 "flat vector index range overflows usize",
798 )
799 })?;
800 if end_idx > self.num_vectors {
801 return Err(io::Error::new(
802 io::ErrorKind::InvalidInput,
803 format!(
804 "flat vector range {start_idx}..{end_idx} exceeds {} vectors",
805 self.num_vectors
806 ),
807 ));
808 }
809
810 if self.locations.is_some() {
811 let len = count.checked_mul(self.vbs).ok_or_else(|| {
812 io::Error::new(io::ErrorKind::InvalidInput, "vector batch size overflow")
813 })?;
814 if count == 0 {
815 return Ok((0..0, 0));
816 }
817 if self.contiguous_rows(start_idx, count) != count {
818 return Err(io::Error::new(
819 io::ErrorKind::InvalidInput,
820 "vector batch is not physically contiguous",
821 ));
822 }
823 let start = self.code_offset(start_idx);
824 let end = start
825 .checked_add(len as u64)
826 .filter(|&end| end <= self.handle.len())
827 .ok_or_else(|| {
828 io::Error::new(
829 io::ErrorKind::InvalidData,
830 "vector location outside ANN payload",
831 )
832 })?;
833 return Ok((start..end, len));
834 }
835 let relative_offset = start_idx.checked_mul(self.vbs).ok_or_else(|| {
836 io::Error::new(
837 io::ErrorKind::InvalidData,
838 "flat vector byte offset overflows usize",
839 )
840 })?;
841 let byte_len = count.checked_mul(self.vbs).ok_or_else(|| {
842 io::Error::new(
843 io::ErrorKind::InvalidInput,
844 "flat vector byte length overflows usize",
845 )
846 })?;
847 let relative_offset = u64::try_from(relative_offset).map_err(|_| {
848 io::Error::new(
849 io::ErrorKind::InvalidData,
850 "flat vector byte offset does not fit in u64",
851 )
852 })?;
853 let byte_len_u64 = u64::try_from(byte_len).map_err(|_| {
854 io::Error::new(
855 io::ErrorKind::InvalidInput,
856 "flat vector byte length does not fit in u64",
857 )
858 })?;
859 let start = self
860 .vectors_offset
861 .checked_add(relative_offset)
862 .ok_or_else(|| {
863 io::Error::new(
864 io::ErrorKind::InvalidData,
865 "flat vector byte offset overflows u64",
866 )
867 })?;
868 let end = start.checked_add(byte_len_u64).ok_or_else(|| {
869 io::Error::new(
870 io::ErrorKind::InvalidData,
871 "flat vector byte range overflows u64",
872 )
873 })?;
874 let vectors_end = self
875 .vectors_offset
876 .checked_add(self.vectors_byte_len)
877 .ok_or_else(|| {
878 io::Error::new(
879 io::ErrorKind::InvalidData,
880 "flat vector payload boundary overflows u64",
881 )
882 })?;
883 if end > vectors_end || end > self.handle.len() {
884 return Err(io::Error::new(
885 io::ErrorKind::InvalidData,
886 format!(
887 "flat vector byte range {start}..{end} exceeds payload boundary {vectors_end}"
888 ),
889 ));
890 }
891 Ok((start..end, byte_len))
892 }
893
894 #[cfg(feature = "native")]
897 pub(crate) fn pin_doc_ids(
898 &mut self,
899 mode: crate::segment::pin::PinMode,
900 remaining: &mut u64,
901 report: &mut crate::segment::pin::PinReport,
902 ) {
903 crate::segment::pin::pin_section(
904 &mut self.doc_ids_bytes,
905 "flat doc_ids",
906 mode,
907 remaining,
908 report,
909 );
910 }
911
912 #[cfg(feature = "native")]
919 pub fn advise_random_access(&self) {
920 if self.locations.is_some() {
921 return;
922 }
923 let Some(vectors_end) = self.vectors_offset.checked_add(self.vectors_byte_len) else {
924 return;
925 };
926 self.handle
927 .madvise_range(self.vectors_offset..vectors_end, libc::MADV_RANDOM);
928 }
929
930 #[cfg(feature = "native")]
937 pub fn prefetch_vectors(&self, sorted_flat_indexes: impl IntoIterator<Item = usize>) {
938 const COALESCE_GAP: u64 = 64 * 1024;
940 let mut ranges = sorted_flat_indexes.into_iter().filter_map(|idx| {
941 self.checked_vector_range(idx, 1)
942 .ok()
943 .map(|(range, _)| range)
944 });
945 let Some(first) = ranges.next() else {
946 return;
947 };
948 let mut run_start = first.start;
949 let mut run_end = first.end;
950 for range in ranges {
951 if range.start >= run_start && range.start <= run_end.saturating_add(COALESCE_GAP) {
952 run_end = run_end.max(range.end);
953 } else {
954 self.handle
955 .madvise_range(run_start..run_end, libc::MADV_WILLNEED);
956 run_start = range.start;
957 run_end = range.end;
958 }
959 }
960 self.handle
961 .madvise_range(run_start..run_end, libc::MADV_WILLNEED);
962 }
963
964 pub async fn read_vector_into(&self, idx: usize, out: &mut [f32]) -> io::Result<()> {
969 if out.len() < self.dim {
970 return Err(io::Error::new(
971 io::ErrorKind::InvalidInput,
972 format!(
973 "flat vector output is too short: need {} floats, got {}",
974 self.dim,
975 out.len()
976 ),
977 ));
978 }
979 let bytes = self.read_vectors_batch(idx, 1).await?;
980 dequantize_raw(bytes.as_slice(), self.quantization, self.dim, out)
981 }
982
983 pub async fn get_vector(&self, idx: usize) -> io::Result<Vec<f32>> {
985 let mut vector = vec![0f32; self.dim];
986 self.read_vector_into(idx, &mut vector).await?;
987 Ok(vector)
988 }
989
990 pub async fn read_vector_raw_into(&self, idx: usize, out: &mut [u8]) -> io::Result<()> {
995 self.read_vector_prefix_raw_into(idx, self.vector_byte_size(), out)
996 .await
997 }
998
999 pub async fn read_vector_prefix_raw_into(
1005 &self,
1006 idx: usize,
1007 prefix_byte_len: usize,
1008 out: &mut [u8],
1009 ) -> io::Result<()> {
1010 let vbs = self.vector_byte_size();
1011 if prefix_byte_len > vbs {
1012 return Err(io::Error::new(
1013 io::ErrorKind::InvalidInput,
1014 format!(
1015 "vector prefix is {prefix_byte_len} bytes, but a vector has only {vbs} bytes"
1016 ),
1017 ));
1018 }
1019 if out.len() < prefix_byte_len {
1020 return Err(io::Error::new(
1021 io::ErrorKind::InvalidInput,
1022 format!(
1023 "vector prefix output is too short: need {prefix_byte_len} bytes, got {}",
1024 out.len()
1025 ),
1026 ));
1027 }
1028 let (full_range, _) = self.checked_vector_range(idx, 1)?;
1029 if prefix_byte_len == 0 {
1030 return Ok(());
1031 }
1032 let prefix_byte_len_u64 = u64::try_from(prefix_byte_len).map_err(|_| {
1033 io::Error::new(
1034 io::ErrorKind::InvalidInput,
1035 "vector prefix length does not fit in u64",
1036 )
1037 })?;
1038 let byte_end = full_range
1039 .start
1040 .checked_add(prefix_byte_len_u64)
1041 .ok_or_else(|| {
1042 io::Error::new(
1043 io::ErrorKind::InvalidData,
1044 "vector byte range overflows u64",
1045 )
1046 })?;
1047 let bytes = self
1048 .handle
1049 .read_bytes_range(full_range.start..byte_end)
1050 .await?;
1051 if bytes.len() != prefix_byte_len {
1052 return Err(io::Error::new(
1053 io::ErrorKind::UnexpectedEof,
1054 format!(
1055 "vector prefix read returned {} bytes, expected {prefix_byte_len}",
1056 bytes.len()
1057 ),
1058 ));
1059 }
1060 out[..prefix_byte_len].copy_from_slice(bytes.as_slice());
1061 Ok(())
1062 }
1063
1064 pub async fn read_vectors_batch(
1070 &self,
1071 start_idx: usize,
1072 count: usize,
1073 ) -> io::Result<OwnedBytes> {
1074 if self.locations.is_some() {
1075 start_idx
1076 .checked_add(count)
1077 .filter(|&end| end <= self.num_vectors)
1078 .ok_or_else(|| {
1079 io::Error::new(
1080 io::ErrorKind::InvalidInput,
1081 "vector batch outside logical range",
1082 )
1083 })?;
1084 let len = count.checked_mul(self.vbs).ok_or_else(|| {
1085 io::Error::new(io::ErrorKind::InvalidInput, "vector batch size overflow")
1086 })?;
1087 if count == 0 {
1088 return Ok(OwnedBytes::empty());
1089 }
1090 let first_count = self.contiguous_rows(start_idx, count);
1091 if first_count < count {
1092 let bytes = self.handle.read_bytes().await?;
1093 return Ok(self.gather_vectors(bytes.as_slice(), start_idx, count, len));
1094 }
1095 }
1096 let (range, expected_len) = self.checked_vector_range(start_idx, count)?;
1097 let bytes = self.handle.read_bytes_range(range).await?;
1098 if bytes.len() != expected_len {
1099 return Err(io::Error::new(
1100 io::ErrorKind::UnexpectedEof,
1101 format!(
1102 "flat vector batch read returned {} bytes, expected {expected_len}",
1103 bytes.len()
1104 ),
1105 ));
1106 }
1107 Ok(bytes)
1108 }
1109
1110 fn gather_vectors(&self, bytes: &[u8], start: usize, count: usize, len: usize) -> OwnedBytes {
1113 let mut gathered = Vec::with_capacity(len);
1114 for row in start..start + count {
1115 let offset = self.code_offset(row) as usize;
1116 gathered.extend_from_slice(&bytes[offset..offset + self.vbs]);
1117 }
1118 OwnedBytes::new(gathered)
1119 }
1120
1121 #[cfg(feature = "sync")]
1123 pub fn read_vector_raw_into_sync(&self, idx: usize, out: &mut [u8]) -> io::Result<()> {
1124 let vbs = self.vector_byte_size();
1125 if out.len() < vbs {
1126 return Err(io::Error::new(
1127 io::ErrorKind::InvalidInput,
1128 format!(
1129 "flat vector output is too short: need {vbs} bytes, got {}",
1130 out.len()
1131 ),
1132 ));
1133 }
1134 let bytes = self.read_vectors_batch_sync(idx, 1)?;
1135 out[..vbs].copy_from_slice(bytes.as_slice());
1136 Ok(())
1137 }
1138
1139 #[cfg(feature = "sync")]
1141 pub fn read_vectors_batch_sync(
1142 &self,
1143 start_idx: usize,
1144 count: usize,
1145 ) -> io::Result<OwnedBytes> {
1146 if self.locations.is_some() {
1147 start_idx
1148 .checked_add(count)
1149 .filter(|&end| end <= self.num_vectors)
1150 .ok_or_else(|| {
1151 io::Error::new(
1152 io::ErrorKind::InvalidInput,
1153 "vector batch outside logical range",
1154 )
1155 })?;
1156 let len = count.checked_mul(self.vbs).ok_or_else(|| {
1157 io::Error::new(io::ErrorKind::InvalidInput, "vector batch size overflow")
1158 })?;
1159 if count == 0 {
1160 return Ok(OwnedBytes::empty());
1161 }
1162 let first_count = self.contiguous_rows(start_idx, count);
1163 if first_count < count {
1164 let bytes = self.handle.read_bytes_sync()?;
1165 return Ok(self.gather_vectors(bytes.as_slice(), start_idx, count, len));
1166 }
1167 }
1168 let (range, expected_len) = self.checked_vector_range(start_idx, count)?;
1169 let bytes = self.handle.read_bytes_range_sync(range)?;
1170 if bytes.len() != expected_len {
1171 return Err(io::Error::new(
1172 io::ErrorKind::UnexpectedEof,
1173 format!(
1174 "flat vector batch read returned {} bytes, expected {expected_len}",
1175 bytes.len()
1176 ),
1177 ));
1178 }
1179 Ok(bytes)
1180 }
1181
1182 pub fn flat_indexes_for_doc_range(&self, doc_id: u32) -> (usize, usize) {
1188 let n = self.num_vectors;
1189 let start = {
1190 let mut lo = 0usize;
1191 let mut hi = n;
1192 while lo < hi {
1193 let mid = lo + (hi - lo) / 2;
1194 if self.doc_id_at(mid) < doc_id {
1195 lo = mid + 1;
1196 } else {
1197 hi = mid;
1198 }
1199 }
1200 lo
1201 };
1202 let mut count = 0;
1203 let mut i = start;
1204 while i < n && self.doc_id_at(i) == doc_id {
1205 count += 1;
1206 i += 1;
1207 }
1208 (start, count)
1209 }
1210
1211 pub fn flat_indexes_for_doc(&self, doc_id: u32) -> (usize, Vec<(u32, u16)>) {
1218 let n = self.num_vectors;
1219 let start = {
1221 let mut lo = 0usize;
1222 let mut hi = n;
1223 while lo < hi {
1224 let mid = lo + (hi - lo) / 2;
1225 if self.doc_id_at(mid) < doc_id {
1226 lo = mid + 1;
1227 } else {
1228 hi = mid;
1229 }
1230 }
1231 lo
1232 };
1233 let mut entries = Vec::new();
1235 let mut i = start;
1236 while i < n {
1237 let (did, ord) = self.get_doc_id(i);
1238 if did != doc_id {
1239 break;
1240 }
1241 entries.push((did, ord));
1242 i += 1;
1243 }
1244 (start, entries)
1245 }
1246
1247 #[inline]
1250 fn doc_map_entry(&self, idx: usize) -> &[u8; DOC_ID_ENTRY_SIZE] {
1251 assert!(
1252 self.has_doc_map,
1253 "document IDs are unavailable on a training-only flat-vector reader",
1254 );
1255 let off = idx * self.doc_map_stride();
1256 self.doc_ids_bytes[off..off + DOC_ID_ENTRY_SIZE]
1257 .try_into()
1258 .expect("doc-map entry slice is DOC_ID_ENTRY_SIZE bytes")
1259 }
1260
1261 #[inline]
1263 fn doc_id_at(&self, idx: usize) -> u32 {
1264 let d = self.doc_map_entry(idx);
1265 u32::from_le_bytes([d[0], d[1], d[2], d[3]])
1266 + self
1267 .locations
1268 .as_ref()
1269 .map_or(0, |layout| layout.doc_base(idx))
1270 }
1271
1272 #[inline]
1274 pub fn get_doc_id(&self, idx: usize) -> (u32, u16) {
1275 let d = self.doc_map_entry(idx);
1276 let doc_id = u32::from_le_bytes([d[0], d[1], d[2], d[3]])
1277 + self
1278 .locations
1279 .as_ref()
1280 .map_or(0, |layout| layout.doc_base(idx));
1281 let ordinal = u16::from_le_bytes([d[4], d[5]]);
1282 (doc_id, ordinal)
1283 }
1284
1285 #[inline]
1287 pub fn vector_byte_size(&self) -> usize {
1288 self.vbs
1289 }
1290
1291 #[inline]
1293 pub fn num_docs_with_vectors(&self) -> usize {
1294 self.num_docs_with_vectors
1295 }
1296
1297 #[cfg(feature = "native")]
1300 pub(crate) fn flat_region(&self) -> Option<(&FileHandle, std::ops::Range<u64>)> {
1301 self.locations.is_none().then(|| {
1302 (
1303 &self.handle,
1304 self.vectors_offset..self.vectors_offset + self.vectors_byte_len,
1305 )
1306 })
1307 }
1308
1309 pub fn vector_bytes_len(&self) -> u64 {
1311 self.vectors_byte_len
1312 }
1313
1314 pub fn exact_lookup_bytes(&self) -> u64 {
1316 self.locations
1317 .as_ref()
1318 .map_or(0, |layout| layout.serialized_len(self.num_vectors))
1319 }
1320
1321 pub fn estimated_heap_bytes(&self) -> usize {
1323 size_of::<Self>()
1324 + self
1325 .locations
1326 .as_ref()
1327 .map_or(0, |layout| layout.heap_bytes())
1328 }
1329}
1330
1331#[cfg(test)]
1332mod tests {
1333 use super::*;
1334
1335 #[test]
1336 fn dequantize_raw_accepts_valid_storage_formats() {
1337 let f32_values = [1.25f32, -2.5];
1338 let f32_bytes = unsafe {
1339 std::slice::from_raw_parts(
1341 f32_values.as_ptr().cast::<u8>(),
1342 std::mem::size_of_val(&f32_values),
1343 )
1344 };
1345 let mut out = [0.0; 2];
1346 dequantize_raw(f32_bytes, DenseVectorQuantization::F32, 2, &mut out).unwrap();
1347 assert_eq!(out, f32_values);
1348
1349 let f16_values = [0x3c00u16, 0xc000u16]; let f16_bytes = unsafe {
1351 std::slice::from_raw_parts(
1353 f16_values.as_ptr().cast::<u8>(),
1354 std::mem::size_of_val(&f16_values),
1355 )
1356 };
1357 dequantize_raw(f16_bytes, DenseVectorQuantization::F16, 2, &mut out).unwrap();
1358 assert_eq!(out, [1.0, -2.0]);
1359
1360 dequantize_raw(&[0, u8::MAX], DenseVectorQuantization::UInt8, 2, &mut out).unwrap();
1361 assert_eq!(out, [u8_to_f32(0), u8_to_f32(u8::MAX)]);
1362 }
1363
1364 #[test]
1365 fn dequantize_raw_rejects_invalid_lengths_and_binary_storage() {
1366 let mut out = [0.0; 2];
1367 assert_eq!(
1368 dequantize_raw(&[0; 7], DenseVectorQuantization::F32, 2, &mut out)
1369 .unwrap_err()
1370 .kind(),
1371 io::ErrorKind::InvalidData
1372 );
1373 assert_eq!(
1374 dequantize_raw(&[0; 8], DenseVectorQuantization::F32, 2, &mut out[..1])
1375 .unwrap_err()
1376 .kind(),
1377 io::ErrorKind::InvalidInput
1378 );
1379 assert_eq!(
1380 dequantize_raw(&[], DenseVectorQuantization::Binary, 0, &mut [])
1381 .unwrap_err()
1382 .kind(),
1383 io::ErrorKind::InvalidInput
1384 );
1385 }
1386
1387 #[test]
1388 fn dequantize_raw_rejects_misaligned_typed_storage() {
1389 let storage = [0u8; 9];
1390 let offset = if (storage.as_ptr() as usize).is_multiple_of(4) {
1391 1
1392 } else {
1393 0
1394 };
1395 let raw = &storage[offset..offset + 8];
1396 assert!(!(raw.as_ptr() as usize).is_multiple_of(4));
1397
1398 let mut out = [0.0; 2];
1399 assert_eq!(
1400 dequantize_raw(raw, DenseVectorQuantization::F32, 2, &mut out)
1401 .unwrap_err()
1402 .kind(),
1403 io::ErrorKind::InvalidData
1404 );
1405 }
1406
1407 #[test]
1408 fn flat_vector_writers_reject_inconsistent_shapes_and_doc_maps() {
1409 let mut encoded = Vec::new();
1410 assert!(
1411 FlatVectorData::serialize_binary_from_flat_streaming(
1412 0,
1413 &[],
1414 &[],
1415 DenseVectorQuantization::F32,
1416 &mut encoded,
1417 )
1418 .is_err()
1419 );
1420 assert!(encoded.is_empty());
1421
1422 assert!(
1423 FlatVectorData::serialize_binary_from_flat_streaming(
1424 2,
1425 &[1.0],
1426 &[(0, 0)],
1427 DenseVectorQuantization::F32,
1428 &mut encoded,
1429 )
1430 .is_err()
1431 );
1432 assert!(encoded.is_empty());
1433
1434 assert!(
1435 FlatVectorData::serialize_binary_from_flat_streaming(
1436 1,
1437 &[1.0],
1438 &[(0, 0)],
1439 DenseVectorQuantization::Binary,
1440 &mut encoded,
1441 )
1442 .is_err()
1443 );
1444 assert!(encoded.is_empty());
1445
1446 assert!(
1447 FlatVectorData::serialize_binary_from_flat_streaming(
1448 1,
1449 &[1.0, 2.0],
1450 &[(1, 0), (0, 0)],
1451 DenseVectorQuantization::F32,
1452 &mut encoded,
1453 )
1454 .is_err()
1455 );
1456 assert!(encoded.is_empty());
1457
1458 assert!(
1459 FlatVectorData::serialize_binary_from_flat_streaming(
1460 1,
1461 &[1.0, 2.0],
1462 &[(0, 0), (0, 0)],
1463 DenseVectorQuantization::F32,
1464 &mut encoded,
1465 )
1466 .is_err()
1467 );
1468 assert!(encoded.is_empty());
1469
1470 assert!(
1471 FlatVectorData::serialize_binary_from_bits_streaming(7, &[0], &[(0, 0)], &mut encoded,)
1472 .is_err()
1473 );
1474 assert!(encoded.is_empty());
1475
1476 assert!(
1477 FlatVectorData::serialize_binary_from_bits_streaming(8, &[], &[(0, 0)], &mut encoded,)
1478 .is_err()
1479 );
1480 assert!(encoded.is_empty());
1481
1482 let vectors = [1.0f32, 2.0, 3.0, 4.0];
1483 let doc_ids = [(0, 0), (1, 0)];
1484 FlatVectorData::serialize_binary_from_flat_streaming(
1485 2,
1486 &vectors,
1487 &doc_ids,
1488 DenseVectorQuantization::F32,
1489 &mut encoded,
1490 )
1491 .unwrap();
1492 assert_eq!(
1493 encoded.len(),
1494 FlatVectorData::serialized_binary_size(2, 2, DenseVectorQuantization::F32).unwrap()
1495 );
1496 }
1497
1498 fn encoded_two_vector_payload() -> Vec<u8> {
1499 let mut encoded = Vec::new();
1500 FlatVectorData::serialize_binary_from_flat_streaming(
1501 2,
1502 &[1.0, 2.0, 3.0, 4.0],
1503 &[(0, 0), (1, 0)],
1504 DenseVectorQuantization::F32,
1505 &mut encoded,
1506 )
1507 .unwrap();
1508 encoded
1509 }
1510
1511 #[tokio::test]
1512 async fn flat_vector_open_rejects_corrupt_layout_and_doc_map() {
1513 let valid = encoded_two_vector_payload();
1514
1515 let mut multi_value = Vec::new();
1516 FlatVectorData::serialize_binary_from_flat_streaming(
1517 1,
1518 &[1.0, 2.0, 3.0],
1519 &[(0, 0), (0, 1), (2, 0)],
1520 DenseVectorQuantization::F32,
1521 &mut multi_value,
1522 )
1523 .unwrap();
1524 let multi_value = LazyFlatVectorData::open_with_doc_limit(
1525 FileHandle::from_bytes(OwnedBytes::new(multi_value)),
1526 Some(3),
1527 )
1528 .await
1529 .unwrap();
1530 assert_eq!(multi_value.num_docs_with_vectors(), 2);
1531
1532 let mut trailing = valid.clone();
1533 trailing.push(0);
1534 assert!(
1535 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(trailing)))
1536 .await
1537 .is_err()
1538 );
1539
1540 let mut truncated = valid.clone();
1541 truncated.pop();
1542 assert!(
1543 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(truncated)))
1544 .await
1545 .is_err()
1546 );
1547
1548 let mut reserved = valid.clone();
1549 reserved[13] = 1;
1550 assert!(
1551 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(reserved)))
1552 .await
1553 .is_err()
1554 );
1555
1556 let doc_map_start = FLAT_BINARY_HEADER_SIZE + 2 * 2 * size_of::<f32>();
1557 let mut unsorted = valid.clone();
1558 let (first, second) = unsorted[doc_map_start..doc_map_start + 2 * DOC_ID_ENTRY_SIZE]
1559 .split_at_mut(DOC_ID_ENTRY_SIZE);
1560 first.swap_with_slice(second);
1561 assert!(
1562 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(unsorted)))
1563 .await
1564 .is_err()
1565 );
1566
1567 let mut duplicate = valid.clone();
1568 duplicate.copy_within(
1569 doc_map_start..doc_map_start + DOC_ID_ENTRY_SIZE,
1570 doc_map_start + DOC_ID_ENTRY_SIZE,
1571 );
1572 assert!(
1573 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(duplicate)))
1574 .await
1575 .is_err()
1576 );
1577
1578 assert!(
1579 LazyFlatVectorData::open_with_doc_limit(
1580 FileHandle::from_bytes(OwnedBytes::new(valid)),
1581 Some(1),
1582 )
1583 .await
1584 .is_err()
1585 );
1586
1587 let mut invalid_binary = Vec::new();
1588 FlatVectorData::serialize_binary_from_bits_streaming(
1589 8,
1590 &[0],
1591 &[(0, 0)],
1592 &mut invalid_binary,
1593 )
1594 .unwrap();
1595 invalid_binary[4..8].copy_from_slice(&7u32.to_le_bytes());
1596 assert!(
1597 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(invalid_binary)))
1598 .await
1599 .is_err()
1600 );
1601 }
1602
1603 #[tokio::test]
1604 async fn flat_vector_batch_and_dequantized_reads_are_checked() {
1605 let flat = LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(
1606 encoded_two_vector_payload(),
1607 )))
1608 .await
1609 .unwrap();
1610
1611 assert_eq!(flat.read_vectors_batch(0, 2).await.unwrap().len(), 16);
1612 assert_eq!(flat.read_vectors_batch(2, 0).await.unwrap().len(), 0);
1613 assert!(flat.read_vectors_batch(1, 2).await.is_err());
1614 assert!(flat.read_vectors_batch(usize::MAX, 1).await.is_err());
1615 assert!(flat.read_vectors_batch(0, usize::MAX).await.is_err());
1616
1617 let mut values = [0.0; 2];
1618 flat.read_vector_into(1, &mut values).await.unwrap();
1619 assert_eq!(values, [3.0, 4.0]);
1620 assert!(flat.read_vector_into(2, &mut values).await.is_err());
1621 assert!(flat.read_vector_into(0, &mut values[..1]).await.is_err());
1622
1623 #[cfg(feature = "sync")]
1624 {
1625 assert_eq!(flat.read_vectors_batch_sync(0, 2).unwrap().len(), 16);
1626 assert!(flat.read_vectors_batch_sync(1, 2).is_err());
1627 assert!(flat.read_vectors_batch_sync(usize::MAX, 1).is_err());
1628 let mut too_short = [0; 7];
1629 assert!(flat.read_vector_raw_into_sync(0, &mut too_short).is_err());
1630 }
1631 }
1632
1633 #[cfg(not(target_arch = "wasm32"))]
1634 #[tokio::test]
1635 async fn flat_vector_reads_reject_short_lazy_range_results() {
1636 let payload = std::sync::Arc::new(encoded_two_vector_payload());
1637 let payload_len = payload.len() as u64;
1638 let read_payload = std::sync::Arc::clone(&payload);
1639 let read_fn: crate::directories::RangeReadFn = std::sync::Arc::new(move |range| {
1640 let payload = std::sync::Arc::clone(&read_payload);
1641 Box::pin(async move {
1642 let start = usize::try_from(range.start).unwrap();
1643 let mut end = usize::try_from(range.end).unwrap();
1644 if range.start == FLAT_BINARY_HEADER_SIZE as u64 {
1647 end -= 1;
1648 }
1649 Ok(OwnedBytes::new(payload[start..end].to_vec()))
1650 })
1651 });
1652 let flat = LazyFlatVectorData::open(FileHandle::lazy(payload_len, read_fn))
1653 .await
1654 .unwrap();
1655
1656 let error = flat.read_vectors_batch(0, 1).await.unwrap_err();
1657 assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
1658 let mut raw = [0; 8];
1659 let error = flat.read_vector_raw_into(0, &mut raw).await.unwrap_err();
1660 assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
1661 }
1662
1663 #[tokio::test]
1664 async fn vector_prefix_reads_are_checked_and_do_not_fetch_the_tail() {
1665 let vectors = [1.0f32, 2.0, 3.0, 4.0];
1666 let doc_ids = [(0, 0), (1, 0)];
1667 let mut encoded = Vec::new();
1668 FlatVectorData::serialize_binary_from_flat_streaming(
1669 2,
1670 &vectors,
1671 &doc_ids,
1672 DenseVectorQuantization::F32,
1673 &mut encoded,
1674 )
1675 .unwrap();
1676 let flat = LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(encoded)))
1677 .await
1678 .unwrap();
1679
1680 let mut prefix = [0xa5; 8];
1681 flat.read_vector_prefix_raw_into(1, 4, &mut prefix)
1682 .await
1683 .unwrap();
1684 assert_eq!(&prefix[..4], &3.0f32.to_ne_bytes());
1685 assert_eq!(&prefix[4..], &[0xa5; 4]);
1686
1687 let mut full = [0; 8];
1688 flat.read_vector_raw_into(1, &mut full).await.unwrap();
1689 assert_eq!(&full[..4], &3.0f32.to_ne_bytes());
1690 assert_eq!(&full[4..], &4.0f32.to_ne_bytes());
1691
1692 assert!(
1693 flat.read_vector_prefix_raw_into(2, 4, &mut prefix)
1694 .await
1695 .is_err()
1696 );
1697 assert!(
1698 flat.read_vector_prefix_raw_into(0, 9, &mut prefix)
1699 .await
1700 .is_err()
1701 );
1702 assert!(
1703 flat.read_vector_prefix_raw_into(0, 4, &mut prefix[..3])
1704 .await
1705 .is_err()
1706 );
1707 }
1708}