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 handle: FileHandle,
406 vectors_offset: u64,
408 vbs: usize,
410 vectors_byte_len: u64,
412}
413
414impl LazyFlatVectorData {
415 pub async fn open(handle: FileHandle) -> io::Result<Self> {
420 Self::open_with_doc_limit(handle, None).await
421 }
422
423 pub(crate) async fn open_with_doc_limit(
429 handle: FileHandle,
430 total_docs: Option<u32>,
431 ) -> io::Result<Self> {
432 Self::open_impl(handle, total_docs, true).await
433 }
434
435 pub(crate) async fn open_for_training(handle: FileHandle) -> io::Result<Self> {
441 Self::open_impl(handle, None, false).await
442 }
443
444 async fn open_impl(
445 handle: FileHandle,
446 total_docs: Option<u32>,
447 load_doc_map: bool,
448 ) -> io::Result<Self> {
449 let header_len = u64::try_from(FLAT_BINARY_HEADER_SIZE).map_err(|_| {
450 io::Error::new(
451 io::ErrorKind::InvalidData,
452 "flat vector header size does not fit in u64",
453 )
454 })?;
455 if handle.len() < header_len {
456 return Err(io::Error::new(
457 io::ErrorKind::UnexpectedEof,
458 format!(
459 "flat vector payload is {} bytes, shorter than its {FLAT_BINARY_HEADER_SIZE}-byte header",
460 handle.len()
461 ),
462 ));
463 }
464
465 let header = handle.read_bytes_range(0..header_len).await?;
467 if header.len() != FLAT_BINARY_HEADER_SIZE {
468 return Err(io::Error::new(
469 io::ErrorKind::UnexpectedEof,
470 format!(
471 "flat vector header read returned {} bytes, expected {FLAT_BINARY_HEADER_SIZE}",
472 header.len()
473 ),
474 ));
475 }
476 let hdr = header.as_slice();
477
478 let magic = u32::from_le_bytes([hdr[0], hdr[1], hdr[2], hdr[3]]);
479 if magic != FLAT_BINARY_MAGIC {
480 return Err(io::Error::new(
481 io::ErrorKind::InvalidData,
482 "Invalid FlatVectorData binary magic",
483 ));
484 }
485
486 let dim = u32::from_le_bytes([hdr[4], hdr[5], hdr[6], hdr[7]]) as usize;
487 let num_vectors = u32::from_le_bytes([hdr[8], hdr[9], hdr[10], hdr[11]]) as usize;
488 let quantization = DenseVectorQuantization::from_tag(hdr[12]).ok_or_else(|| {
489 io::Error::new(
490 io::ErrorKind::InvalidData,
491 format!("Unknown quantization tag: {}", hdr[12]),
492 )
493 })?;
494 if hdr[13..] != [0, 0, 0] {
495 return Err(io::Error::new(
496 io::ErrorKind::InvalidData,
497 "flat vector header has non-zero reserved bytes",
498 ));
499 }
500
501 let vbs =
503 FlatVectorData::validate_shape(dim, num_vectors, quantization).map_err(|error| {
504 io::Error::new(
505 io::ErrorKind::InvalidData,
506 format!("invalid flat vector shape: {error}"),
507 )
508 })?;
509 let vectors_byte_len_usize = num_vectors.checked_mul(vbs).ok_or_else(|| {
510 io::Error::new(
511 io::ErrorKind::InvalidData,
512 "flat vector payload size overflows usize",
513 )
514 })?;
515 let doc_ids_byte_len_usize =
516 num_vectors.checked_mul(DOC_ID_ENTRY_SIZE).ok_or_else(|| {
517 io::Error::new(
518 io::ErrorKind::InvalidData,
519 "flat vector doc-map size overflows usize",
520 )
521 })?;
522 let expected_len_usize = FLAT_BINARY_HEADER_SIZE
523 .checked_add(vectors_byte_len_usize)
524 .and_then(|size| size.checked_add(doc_ids_byte_len_usize))
525 .ok_or_else(|| {
526 io::Error::new(
527 io::ErrorKind::InvalidData,
528 "flat vector serialized size overflows usize",
529 )
530 })?;
531 let expected_len = u64::try_from(expected_len_usize).map_err(|_| {
532 io::Error::new(
533 io::ErrorKind::InvalidData,
534 "flat vector serialized size does not fit in u64",
535 )
536 })?;
537 if handle.len() != expected_len {
538 return Err(io::Error::new(
539 io::ErrorKind::InvalidData,
540 format!(
541 "flat vector payload has {} bytes, expected exactly {expected_len}",
542 handle.len()
543 ),
544 ));
545 }
546
547 let vectors_byte_len = u64::try_from(vectors_byte_len_usize).map_err(|_| {
548 io::Error::new(
549 io::ErrorKind::InvalidData,
550 "flat vector payload size does not fit in u64",
551 )
552 })?;
553 let doc_ids_byte_len = u64::try_from(doc_ids_byte_len_usize).map_err(|_| {
554 io::Error::new(
555 io::ErrorKind::InvalidData,
556 "flat vector doc-map size does not fit in u64",
557 )
558 })?;
559 let doc_ids_start = header_len.checked_add(vectors_byte_len).ok_or_else(|| {
560 io::Error::new(
561 io::ErrorKind::InvalidData,
562 "flat vector doc-map offset overflows u64",
563 )
564 })?;
565 let doc_ids_end = doc_ids_start.checked_add(doc_ids_byte_len).ok_or_else(|| {
566 io::Error::new(
567 io::ErrorKind::InvalidData,
568 "flat vector doc-map range overflows u64",
569 )
570 })?;
571
572 let doc_ids_bytes = if load_doc_map {
573 let bytes = handle.read_bytes_range(doc_ids_start..doc_ids_end).await?;
574 if bytes.len() != doc_ids_byte_len_usize {
575 return Err(io::Error::new(
576 io::ErrorKind::UnexpectedEof,
577 format!(
578 "flat vector doc-map read returned {} bytes, expected {doc_ids_byte_len_usize}",
579 bytes.len()
580 ),
581 ));
582 }
583 bytes
584 } else {
585 OwnedBytes::empty()
586 };
587
588 let mut previous = None;
589 let mut num_docs_with_vectors = 0usize;
590 for entry in doc_ids_bytes.as_slice().chunks_exact(DOC_ID_ENTRY_SIZE) {
591 let doc_id = u32::from_le_bytes([entry[0], entry[1], entry[2], entry[3]]);
592 let ordinal = u16::from_le_bytes([entry[4], entry[5]]);
593 let current = (doc_id, ordinal);
594 if let Some(previous) = previous
595 && previous >= current
596 {
597 return Err(io::Error::new(
598 io::ErrorKind::InvalidData,
599 format!(
600 "flat vector doc map must be strictly sorted by (doc_id, ordinal), found {previous:?} before {current:?}"
601 ),
602 ));
603 }
604 if let Some(limit) = total_docs
605 && doc_id >= limit
606 {
607 return Err(io::Error::new(
608 io::ErrorKind::InvalidData,
609 format!(
610 "flat vector doc map references document {doc_id}, but segment contains only {} documents",
611 limit
612 ),
613 ));
614 }
615 if previous.is_none_or(|(previous_doc_id, _)| previous_doc_id != doc_id) {
616 num_docs_with_vectors = num_docs_with_vectors.checked_add(1).ok_or_else(|| {
617 io::Error::new(
618 io::ErrorKind::InvalidData,
619 "flat vector distinct-document count overflows usize",
620 )
621 })?;
622 }
623 previous = Some(current);
624 }
625
626 Ok(Self {
627 dim,
628 num_vectors,
629 num_docs_with_vectors,
630 quantization,
631 doc_ids_bytes,
632 handle,
633 vectors_offset: header_len,
634 vbs,
635 vectors_byte_len,
636 })
637 }
638
639 fn checked_vector_range(
640 &self,
641 start_idx: usize,
642 count: usize,
643 ) -> io::Result<(std::ops::Range<u64>, usize)> {
644 let end_idx = start_idx.checked_add(count).ok_or_else(|| {
645 io::Error::new(
646 io::ErrorKind::InvalidInput,
647 "flat vector index range overflows usize",
648 )
649 })?;
650 if end_idx > self.num_vectors {
651 return Err(io::Error::new(
652 io::ErrorKind::InvalidInput,
653 format!(
654 "flat vector range {start_idx}..{end_idx} exceeds {} vectors",
655 self.num_vectors
656 ),
657 ));
658 }
659
660 let relative_offset = start_idx.checked_mul(self.vbs).ok_or_else(|| {
661 io::Error::new(
662 io::ErrorKind::InvalidData,
663 "flat vector byte offset overflows usize",
664 )
665 })?;
666 let byte_len = count.checked_mul(self.vbs).ok_or_else(|| {
667 io::Error::new(
668 io::ErrorKind::InvalidInput,
669 "flat vector byte length overflows usize",
670 )
671 })?;
672 let relative_offset = u64::try_from(relative_offset).map_err(|_| {
673 io::Error::new(
674 io::ErrorKind::InvalidData,
675 "flat vector byte offset does not fit in u64",
676 )
677 })?;
678 let byte_len_u64 = u64::try_from(byte_len).map_err(|_| {
679 io::Error::new(
680 io::ErrorKind::InvalidInput,
681 "flat vector byte length does not fit in u64",
682 )
683 })?;
684 let start = self
685 .vectors_offset
686 .checked_add(relative_offset)
687 .ok_or_else(|| {
688 io::Error::new(
689 io::ErrorKind::InvalidData,
690 "flat vector byte offset overflows u64",
691 )
692 })?;
693 let end = start.checked_add(byte_len_u64).ok_or_else(|| {
694 io::Error::new(
695 io::ErrorKind::InvalidData,
696 "flat vector byte range overflows u64",
697 )
698 })?;
699 let vectors_end = self
700 .vectors_offset
701 .checked_add(self.vectors_byte_len)
702 .ok_or_else(|| {
703 io::Error::new(
704 io::ErrorKind::InvalidData,
705 "flat vector payload boundary overflows u64",
706 )
707 })?;
708 if end > vectors_end || end > self.handle.len() {
709 return Err(io::Error::new(
710 io::ErrorKind::InvalidData,
711 format!(
712 "flat vector byte range {start}..{end} exceeds payload boundary {vectors_end}"
713 ),
714 ));
715 }
716 Ok((start..end, byte_len))
717 }
718
719 #[cfg(feature = "native")]
722 pub(crate) fn pin_doc_ids(
723 &mut self,
724 mode: crate::segment::pin::PinMode,
725 remaining: &mut u64,
726 report: &mut crate::segment::pin::PinReport,
727 ) {
728 crate::segment::pin::pin_section(
729 &mut self.doc_ids_bytes,
730 "flat doc_ids",
731 mode,
732 remaining,
733 report,
734 );
735 }
736
737 #[cfg(feature = "native")]
744 pub fn advise_random_access(&self) {
745 let Some(vectors_end) = self.vectors_offset.checked_add(self.vectors_byte_len) else {
746 return;
747 };
748 self.handle
749 .madvise_range(self.vectors_offset..vectors_end, libc::MADV_RANDOM);
750 }
751
752 #[cfg(feature = "native")]
759 pub fn prefetch_vectors(&self, sorted_flat_indexes: impl IntoIterator<Item = usize>) {
760 const COALESCE_GAP: u64 = 64 * 1024;
762 let mut ranges = sorted_flat_indexes.into_iter().filter_map(|idx| {
763 self.checked_vector_range(idx, 1)
764 .ok()
765 .map(|(range, _)| range)
766 });
767 let Some(first) = ranges.next() else {
768 return;
769 };
770 let mut run_start = first.start;
771 let mut run_end = first.end;
772 for range in ranges {
773 if range.start <= run_end.saturating_add(COALESCE_GAP) {
774 run_end = run_end.max(range.end);
775 } else {
776 self.handle
777 .madvise_range(run_start..run_end, libc::MADV_WILLNEED);
778 run_start = range.start;
779 run_end = range.end;
780 }
781 }
782 self.handle
783 .madvise_range(run_start..run_end, libc::MADV_WILLNEED);
784 }
785
786 pub async fn read_vector_into(&self, idx: usize, out: &mut [f32]) -> io::Result<()> {
791 if out.len() < self.dim {
792 return Err(io::Error::new(
793 io::ErrorKind::InvalidInput,
794 format!(
795 "flat vector output is too short: need {} floats, got {}",
796 self.dim,
797 out.len()
798 ),
799 ));
800 }
801 let bytes = self.read_vectors_batch(idx, 1).await?;
802 dequantize_raw(bytes.as_slice(), self.quantization, self.dim, out)
803 }
804
805 pub async fn get_vector(&self, idx: usize) -> io::Result<Vec<f32>> {
807 let mut vector = vec![0f32; self.dim];
808 self.read_vector_into(idx, &mut vector).await?;
809 Ok(vector)
810 }
811
812 pub async fn read_vector_raw_into(&self, idx: usize, out: &mut [u8]) -> io::Result<()> {
817 self.read_vector_prefix_raw_into(idx, self.vector_byte_size(), out)
818 .await
819 }
820
821 pub async fn read_vector_prefix_raw_into(
827 &self,
828 idx: usize,
829 prefix_byte_len: usize,
830 out: &mut [u8],
831 ) -> io::Result<()> {
832 let vbs = self.vector_byte_size();
833 if prefix_byte_len > vbs {
834 return Err(io::Error::new(
835 io::ErrorKind::InvalidInput,
836 format!(
837 "vector prefix is {prefix_byte_len} bytes, but a vector has only {vbs} bytes"
838 ),
839 ));
840 }
841 if out.len() < prefix_byte_len {
842 return Err(io::Error::new(
843 io::ErrorKind::InvalidInput,
844 format!(
845 "vector prefix output is too short: need {prefix_byte_len} bytes, got {}",
846 out.len()
847 ),
848 ));
849 }
850 let (full_range, _) = self.checked_vector_range(idx, 1)?;
851 if prefix_byte_len == 0 {
852 return Ok(());
853 }
854 let prefix_byte_len_u64 = u64::try_from(prefix_byte_len).map_err(|_| {
855 io::Error::new(
856 io::ErrorKind::InvalidInput,
857 "vector prefix length does not fit in u64",
858 )
859 })?;
860 let byte_end = full_range
861 .start
862 .checked_add(prefix_byte_len_u64)
863 .ok_or_else(|| {
864 io::Error::new(
865 io::ErrorKind::InvalidData,
866 "vector byte range overflows u64",
867 )
868 })?;
869 let bytes = self
870 .handle
871 .read_bytes_range(full_range.start..byte_end)
872 .await?;
873 if bytes.len() != prefix_byte_len {
874 return Err(io::Error::new(
875 io::ErrorKind::UnexpectedEof,
876 format!(
877 "vector prefix read returned {} bytes, expected {prefix_byte_len}",
878 bytes.len()
879 ),
880 ));
881 }
882 out[..prefix_byte_len].copy_from_slice(bytes.as_slice());
883 Ok(())
884 }
885
886 pub async fn read_vectors_batch(
892 &self,
893 start_idx: usize,
894 count: usize,
895 ) -> io::Result<OwnedBytes> {
896 let (range, expected_len) = self.checked_vector_range(start_idx, count)?;
897 let bytes = self.handle.read_bytes_range(range).await?;
898 if bytes.len() != expected_len {
899 return Err(io::Error::new(
900 io::ErrorKind::UnexpectedEof,
901 format!(
902 "flat vector batch read returned {} bytes, expected {expected_len}",
903 bytes.len()
904 ),
905 ));
906 }
907 Ok(bytes)
908 }
909
910 #[cfg(feature = "sync")]
912 pub fn read_vector_raw_into_sync(&self, idx: usize, out: &mut [u8]) -> io::Result<()> {
913 let vbs = self.vector_byte_size();
914 if out.len() < vbs {
915 return Err(io::Error::new(
916 io::ErrorKind::InvalidInput,
917 format!(
918 "flat vector output is too short: need {vbs} bytes, got {}",
919 out.len()
920 ),
921 ));
922 }
923 let bytes = self.read_vectors_batch_sync(idx, 1)?;
924 out[..vbs].copy_from_slice(bytes.as_slice());
925 Ok(())
926 }
927
928 #[cfg(feature = "sync")]
930 pub fn read_vectors_batch_sync(
931 &self,
932 start_idx: usize,
933 count: usize,
934 ) -> io::Result<OwnedBytes> {
935 let (range, expected_len) = self.checked_vector_range(start_idx, count)?;
936 let bytes = self.handle.read_bytes_range_sync(range)?;
937 if bytes.len() != expected_len {
938 return Err(io::Error::new(
939 io::ErrorKind::UnexpectedEof,
940 format!(
941 "flat vector batch read returned {} bytes, expected {expected_len}",
942 bytes.len()
943 ),
944 ));
945 }
946 Ok(bytes)
947 }
948
949 pub fn flat_indexes_for_doc_range(&self, doc_id: u32) -> (usize, usize) {
955 let n = self.num_vectors;
956 let start = {
957 let mut lo = 0usize;
958 let mut hi = n;
959 while lo < hi {
960 let mid = lo + (hi - lo) / 2;
961 if self.doc_id_at(mid) < doc_id {
962 lo = mid + 1;
963 } else {
964 hi = mid;
965 }
966 }
967 lo
968 };
969 let mut count = 0;
970 let mut i = start;
971 while i < n && self.doc_id_at(i) == doc_id {
972 count += 1;
973 i += 1;
974 }
975 (start, count)
976 }
977
978 pub fn flat_indexes_for_doc(&self, doc_id: u32) -> (usize, Vec<(u32, u16)>) {
985 let n = self.num_vectors;
986 let start = {
988 let mut lo = 0usize;
989 let mut hi = n;
990 while lo < hi {
991 let mid = lo + (hi - lo) / 2;
992 if self.doc_id_at(mid) < doc_id {
993 lo = mid + 1;
994 } else {
995 hi = mid;
996 }
997 }
998 lo
999 };
1000 let mut entries = Vec::new();
1002 let mut i = start;
1003 while i < n {
1004 let (did, ord) = self.get_doc_id(i);
1005 if did != doc_id {
1006 break;
1007 }
1008 entries.push((did, ord));
1009 i += 1;
1010 }
1011 (start, entries)
1012 }
1013
1014 #[inline]
1016 fn doc_id_at(&self, idx: usize) -> u32 {
1017 assert_eq!(
1018 self.doc_ids_bytes.len(),
1019 self.num_vectors * DOC_ID_ENTRY_SIZE,
1020 "document IDs are unavailable on a training-only flat-vector reader",
1021 );
1022 let off = idx * DOC_ID_ENTRY_SIZE;
1023 let d = &self.doc_ids_bytes[off..];
1024 u32::from_le_bytes([d[0], d[1], d[2], d[3]])
1025 }
1026
1027 #[inline]
1029 pub fn get_doc_id(&self, idx: usize) -> (u32, u16) {
1030 assert_eq!(
1031 self.doc_ids_bytes.len(),
1032 self.num_vectors * DOC_ID_ENTRY_SIZE,
1033 "document IDs are unavailable on a training-only flat-vector reader",
1034 );
1035 let off = idx * DOC_ID_ENTRY_SIZE;
1036 let d = &self.doc_ids_bytes[off..];
1037 let doc_id = u32::from_le_bytes([d[0], d[1], d[2], d[3]]);
1038 let ordinal = u16::from_le_bytes([d[4], d[5]]);
1039 (doc_id, ordinal)
1040 }
1041
1042 #[inline]
1044 pub fn vector_byte_size(&self) -> usize {
1045 self.vbs
1046 }
1047
1048 #[inline]
1050 pub fn num_docs_with_vectors(&self) -> usize {
1051 self.num_docs_with_vectors
1052 }
1053
1054 pub fn vector_bytes_len(&self) -> u64 {
1056 self.vectors_byte_len
1057 }
1058
1059 pub fn vectors_byte_offset(&self) -> u64 {
1061 self.vectors_offset
1062 }
1063
1064 pub fn handle(&self) -> &FileHandle {
1066 &self.handle
1067 }
1068
1069 pub fn estimated_memory_bytes(&self) -> usize {
1071 size_of::<Self>() + size_of::<OwnedBytes>()
1072 }
1073}
1074
1075#[cfg(test)]
1076mod tests {
1077 use super::*;
1078
1079 #[test]
1080 fn dequantize_raw_accepts_valid_storage_formats() {
1081 let f32_values = [1.25f32, -2.5];
1082 let f32_bytes = unsafe {
1083 std::slice::from_raw_parts(
1085 f32_values.as_ptr().cast::<u8>(),
1086 std::mem::size_of_val(&f32_values),
1087 )
1088 };
1089 let mut out = [0.0; 2];
1090 dequantize_raw(f32_bytes, DenseVectorQuantization::F32, 2, &mut out).unwrap();
1091 assert_eq!(out, f32_values);
1092
1093 let f16_values = [0x3c00u16, 0xc000u16]; let f16_bytes = unsafe {
1095 std::slice::from_raw_parts(
1097 f16_values.as_ptr().cast::<u8>(),
1098 std::mem::size_of_val(&f16_values),
1099 )
1100 };
1101 dequantize_raw(f16_bytes, DenseVectorQuantization::F16, 2, &mut out).unwrap();
1102 assert_eq!(out, [1.0, -2.0]);
1103
1104 dequantize_raw(&[0, u8::MAX], DenseVectorQuantization::UInt8, 2, &mut out).unwrap();
1105 assert_eq!(out, [u8_to_f32(0), u8_to_f32(u8::MAX)]);
1106 }
1107
1108 #[test]
1109 fn dequantize_raw_rejects_invalid_lengths_and_binary_storage() {
1110 let mut out = [0.0; 2];
1111 assert_eq!(
1112 dequantize_raw(&[0; 7], DenseVectorQuantization::F32, 2, &mut out)
1113 .unwrap_err()
1114 .kind(),
1115 io::ErrorKind::InvalidData
1116 );
1117 assert_eq!(
1118 dequantize_raw(&[0; 8], DenseVectorQuantization::F32, 2, &mut out[..1])
1119 .unwrap_err()
1120 .kind(),
1121 io::ErrorKind::InvalidInput
1122 );
1123 assert_eq!(
1124 dequantize_raw(&[], DenseVectorQuantization::Binary, 0, &mut [])
1125 .unwrap_err()
1126 .kind(),
1127 io::ErrorKind::InvalidInput
1128 );
1129 }
1130
1131 #[test]
1132 fn dequantize_raw_rejects_misaligned_typed_storage() {
1133 let storage = [0u8; 9];
1134 let offset = if (storage.as_ptr() as usize).is_multiple_of(4) {
1135 1
1136 } else {
1137 0
1138 };
1139 let raw = &storage[offset..offset + 8];
1140 assert!(!(raw.as_ptr() as usize).is_multiple_of(4));
1141
1142 let mut out = [0.0; 2];
1143 assert_eq!(
1144 dequantize_raw(raw, DenseVectorQuantization::F32, 2, &mut out)
1145 .unwrap_err()
1146 .kind(),
1147 io::ErrorKind::InvalidData
1148 );
1149 }
1150
1151 #[test]
1152 fn flat_vector_writers_reject_inconsistent_shapes_and_doc_maps() {
1153 let mut encoded = Vec::new();
1154 assert!(
1155 FlatVectorData::serialize_binary_from_flat_streaming(
1156 0,
1157 &[],
1158 &[],
1159 DenseVectorQuantization::F32,
1160 &mut encoded,
1161 )
1162 .is_err()
1163 );
1164 assert!(encoded.is_empty());
1165
1166 assert!(
1167 FlatVectorData::serialize_binary_from_flat_streaming(
1168 2,
1169 &[1.0],
1170 &[(0, 0)],
1171 DenseVectorQuantization::F32,
1172 &mut encoded,
1173 )
1174 .is_err()
1175 );
1176 assert!(encoded.is_empty());
1177
1178 assert!(
1179 FlatVectorData::serialize_binary_from_flat_streaming(
1180 1,
1181 &[1.0],
1182 &[(0, 0)],
1183 DenseVectorQuantization::Binary,
1184 &mut encoded,
1185 )
1186 .is_err()
1187 );
1188 assert!(encoded.is_empty());
1189
1190 assert!(
1191 FlatVectorData::serialize_binary_from_flat_streaming(
1192 1,
1193 &[1.0, 2.0],
1194 &[(1, 0), (0, 0)],
1195 DenseVectorQuantization::F32,
1196 &mut encoded,
1197 )
1198 .is_err()
1199 );
1200 assert!(encoded.is_empty());
1201
1202 assert!(
1203 FlatVectorData::serialize_binary_from_flat_streaming(
1204 1,
1205 &[1.0, 2.0],
1206 &[(0, 0), (0, 0)],
1207 DenseVectorQuantization::F32,
1208 &mut encoded,
1209 )
1210 .is_err()
1211 );
1212 assert!(encoded.is_empty());
1213
1214 assert!(
1215 FlatVectorData::serialize_binary_from_bits_streaming(7, &[0], &[(0, 0)], &mut encoded,)
1216 .is_err()
1217 );
1218 assert!(encoded.is_empty());
1219
1220 assert!(
1221 FlatVectorData::serialize_binary_from_bits_streaming(8, &[], &[(0, 0)], &mut encoded,)
1222 .is_err()
1223 );
1224 assert!(encoded.is_empty());
1225
1226 let vectors = [1.0f32, 2.0, 3.0, 4.0];
1227 let doc_ids = [(0, 0), (1, 0)];
1228 FlatVectorData::serialize_binary_from_flat_streaming(
1229 2,
1230 &vectors,
1231 &doc_ids,
1232 DenseVectorQuantization::F32,
1233 &mut encoded,
1234 )
1235 .unwrap();
1236 assert_eq!(
1237 encoded.len(),
1238 FlatVectorData::serialized_binary_size(2, 2, DenseVectorQuantization::F32).unwrap()
1239 );
1240 }
1241
1242 fn encoded_two_vector_payload() -> Vec<u8> {
1243 let mut encoded = Vec::new();
1244 FlatVectorData::serialize_binary_from_flat_streaming(
1245 2,
1246 &[1.0, 2.0, 3.0, 4.0],
1247 &[(0, 0), (1, 0)],
1248 DenseVectorQuantization::F32,
1249 &mut encoded,
1250 )
1251 .unwrap();
1252 encoded
1253 }
1254
1255 #[tokio::test]
1256 async fn flat_vector_open_rejects_corrupt_layout_and_doc_map() {
1257 let valid = encoded_two_vector_payload();
1258
1259 let mut multi_value = Vec::new();
1260 FlatVectorData::serialize_binary_from_flat_streaming(
1261 1,
1262 &[1.0, 2.0, 3.0],
1263 &[(0, 0), (0, 1), (2, 0)],
1264 DenseVectorQuantization::F32,
1265 &mut multi_value,
1266 )
1267 .unwrap();
1268 let multi_value = LazyFlatVectorData::open_with_doc_limit(
1269 FileHandle::from_bytes(OwnedBytes::new(multi_value)),
1270 Some(3),
1271 )
1272 .await
1273 .unwrap();
1274 assert_eq!(multi_value.num_docs_with_vectors(), 2);
1275
1276 let mut trailing = valid.clone();
1277 trailing.push(0);
1278 assert!(
1279 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(trailing)))
1280 .await
1281 .is_err()
1282 );
1283
1284 let mut truncated = valid.clone();
1285 truncated.pop();
1286 assert!(
1287 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(truncated)))
1288 .await
1289 .is_err()
1290 );
1291
1292 let mut reserved = valid.clone();
1293 reserved[13] = 1;
1294 assert!(
1295 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(reserved)))
1296 .await
1297 .is_err()
1298 );
1299
1300 let doc_map_start = FLAT_BINARY_HEADER_SIZE + 2 * 2 * size_of::<f32>();
1301 let mut unsorted = valid.clone();
1302 let (first, second) = unsorted[doc_map_start..doc_map_start + 2 * DOC_ID_ENTRY_SIZE]
1303 .split_at_mut(DOC_ID_ENTRY_SIZE);
1304 first.swap_with_slice(second);
1305 assert!(
1306 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(unsorted)))
1307 .await
1308 .is_err()
1309 );
1310
1311 let mut duplicate = valid.clone();
1312 duplicate.copy_within(
1313 doc_map_start..doc_map_start + DOC_ID_ENTRY_SIZE,
1314 doc_map_start + DOC_ID_ENTRY_SIZE,
1315 );
1316 assert!(
1317 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(duplicate)))
1318 .await
1319 .is_err()
1320 );
1321
1322 assert!(
1323 LazyFlatVectorData::open_with_doc_limit(
1324 FileHandle::from_bytes(OwnedBytes::new(valid)),
1325 Some(1),
1326 )
1327 .await
1328 .is_err()
1329 );
1330
1331 let mut invalid_binary = Vec::new();
1332 FlatVectorData::serialize_binary_from_bits_streaming(
1333 8,
1334 &[0],
1335 &[(0, 0)],
1336 &mut invalid_binary,
1337 )
1338 .unwrap();
1339 invalid_binary[4..8].copy_from_slice(&7u32.to_le_bytes());
1340 assert!(
1341 LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(invalid_binary)))
1342 .await
1343 .is_err()
1344 );
1345 }
1346
1347 #[tokio::test]
1348 async fn flat_vector_batch_and_dequantized_reads_are_checked() {
1349 let flat = LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(
1350 encoded_two_vector_payload(),
1351 )))
1352 .await
1353 .unwrap();
1354
1355 assert_eq!(flat.read_vectors_batch(0, 2).await.unwrap().len(), 16);
1356 assert_eq!(flat.read_vectors_batch(2, 0).await.unwrap().len(), 0);
1357 assert!(flat.read_vectors_batch(1, 2).await.is_err());
1358 assert!(flat.read_vectors_batch(usize::MAX, 1).await.is_err());
1359 assert!(flat.read_vectors_batch(0, usize::MAX).await.is_err());
1360
1361 let mut values = [0.0; 2];
1362 flat.read_vector_into(1, &mut values).await.unwrap();
1363 assert_eq!(values, [3.0, 4.0]);
1364 assert!(flat.read_vector_into(2, &mut values).await.is_err());
1365 assert!(flat.read_vector_into(0, &mut values[..1]).await.is_err());
1366
1367 #[cfg(feature = "sync")]
1368 {
1369 assert_eq!(flat.read_vectors_batch_sync(0, 2).unwrap().len(), 16);
1370 assert!(flat.read_vectors_batch_sync(1, 2).is_err());
1371 assert!(flat.read_vectors_batch_sync(usize::MAX, 1).is_err());
1372 let mut too_short = [0; 7];
1373 assert!(flat.read_vector_raw_into_sync(0, &mut too_short).is_err());
1374 }
1375 }
1376
1377 #[cfg(not(target_arch = "wasm32"))]
1378 #[tokio::test]
1379 async fn flat_vector_reads_reject_short_lazy_range_results() {
1380 let payload = std::sync::Arc::new(encoded_two_vector_payload());
1381 let payload_len = payload.len() as u64;
1382 let read_payload = std::sync::Arc::clone(&payload);
1383 let read_fn: crate::directories::RangeReadFn = std::sync::Arc::new(move |range| {
1384 let payload = std::sync::Arc::clone(&read_payload);
1385 Box::pin(async move {
1386 let start = usize::try_from(range.start).unwrap();
1387 let mut end = usize::try_from(range.end).unwrap();
1388 if range.start == FLAT_BINARY_HEADER_SIZE as u64 {
1391 end -= 1;
1392 }
1393 Ok(OwnedBytes::new(payload[start..end].to_vec()))
1394 })
1395 });
1396 let flat = LazyFlatVectorData::open(FileHandle::lazy(payload_len, read_fn))
1397 .await
1398 .unwrap();
1399
1400 let error = flat.read_vectors_batch(0, 1).await.unwrap_err();
1401 assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
1402 let mut raw = [0; 8];
1403 let error = flat.read_vector_raw_into(0, &mut raw).await.unwrap_err();
1404 assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
1405 }
1406
1407 #[tokio::test]
1408 async fn vector_prefix_reads_are_checked_and_do_not_fetch_the_tail() {
1409 let vectors = [1.0f32, 2.0, 3.0, 4.0];
1410 let doc_ids = [(0, 0), (1, 0)];
1411 let mut encoded = Vec::new();
1412 FlatVectorData::serialize_binary_from_flat_streaming(
1413 2,
1414 &vectors,
1415 &doc_ids,
1416 DenseVectorQuantization::F32,
1417 &mut encoded,
1418 )
1419 .unwrap();
1420 let flat = LazyFlatVectorData::open(FileHandle::from_bytes(OwnedBytes::new(encoded)))
1421 .await
1422 .unwrap();
1423
1424 let mut prefix = [0xa5; 8];
1425 flat.read_vector_prefix_raw_into(1, 4, &mut prefix)
1426 .await
1427 .unwrap();
1428 assert_eq!(&prefix[..4], &3.0f32.to_ne_bytes());
1429 assert_eq!(&prefix[4..], &[0xa5; 4]);
1430
1431 let mut full = [0; 8];
1432 flat.read_vector_raw_into(1, &mut full).await.unwrap();
1433 assert_eq!(&full[..4], &3.0f32.to_ne_bytes());
1434 assert_eq!(&full[4..], &4.0f32.to_ne_bytes());
1435
1436 assert!(
1437 flat.read_vector_prefix_raw_into(2, 4, &mut prefix)
1438 .await
1439 .is_err()
1440 );
1441 assert!(
1442 flat.read_vector_prefix_raw_into(0, 9, &mut prefix)
1443 .await
1444 .is_err()
1445 );
1446 assert!(
1447 flat.read_vector_prefix_raw_into(0, 4, &mut prefix[..3])
1448 .await
1449 .is_err()
1450 );
1451 }
1452}