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;
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 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 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 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])?; Ok(())
257 }
258
259 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 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 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 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#[derive(Debug, Clone)]
393pub struct LazyFlatVectorData {
394 pub dim: usize,
396 pub num_vectors: usize,
398 num_docs_with_vectors: usize,
400 pub quantization: DenseVectorQuantization,
402 doc_ids_bytes: OwnedBytes,
404 has_doc_map: bool,
408 handle: FileHandle,
410 vectors_offset: u64,
412 vbs: usize,
414 vectors_byte_len: u64,
416}
417
418impl LazyFlatVectorData {
419 pub async fn open(handle: FileHandle) -> io::Result<Self> {
424 Self::open_with_doc_limit(handle, None).await
425 }
426
427 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 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 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 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 #[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 #[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 #[cfg(feature = "native")]
765 pub fn prefetch_vectors(&self, sorted_flat_indexes: impl IntoIterator<Item = usize>) {
766 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 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 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 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 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 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 #[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 #[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 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 pub fn flat_indexes_for_doc(&self, doc_id: u32) -> (usize, Vec<(u32, u16)>) {
991 let n = self.num_vectors;
992 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 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 #[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 #[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 #[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 #[inline]
1052 pub fn vector_byte_size(&self) -> usize {
1053 self.vbs
1054 }
1055
1056 #[inline]
1058 pub fn num_docs_with_vectors(&self) -> usize {
1059 self.num_docs_with_vectors
1060 }
1061
1062 pub fn vector_bytes_len(&self) -> u64 {
1064 self.vectors_byte_len
1065 }
1066
1067 pub fn vectors_byte_offset(&self) -> u64 {
1069 self.vectors_offset
1070 }
1071
1072 pub fn handle(&self) -> &FileHandle {
1074 &self.handle
1075 }
1076
1077 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 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]; let f16_bytes = unsafe {
1103 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 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}