1use std::ptr::NonNull;
9
10use diskann_utils::views::rowmajor;
11
12use super::super::MinMaxQuantizer;
13use super::super::vectors::DataMutRef;
14use crate::CompressInto;
15use crate::bits::{Representation, Unsigned};
16use crate::minmax::{self, Data, MinMaxCompensation};
17use crate::multi_vector::matrix::{
18 Defaulted, NewMut, NewOwned, NewRef, Repr, ReprMut, ReprOwned, SliceError,
19};
20use crate::multi_vector::{LayoutError, Mat, MatMut, MatRef, Overflow, Standard};
21use crate::scalar::InputContainsNaN;
22use crate::utils;
23
24#[derive(Debug, Clone, Copy)]
34pub struct MinMaxMeta<const NBITS: usize> {
35 nrows: usize,
36 intrinsic_dim: usize,
37}
38
39impl<const NBITS: usize> MinMaxMeta<NBITS>
40where
41 Unsigned: Representation<NBITS>,
42{
43 pub fn new(nrows: usize, intrinsic_dim: usize) -> Self {
45 Self {
46 nrows,
47 intrinsic_dim,
48 }
49 }
50
51 pub fn intrinsic_dim(&self) -> usize {
53 self.intrinsic_dim
54 }
55
56 pub fn ncols(&self) -> usize {
59 Data::<NBITS>::canonical_bytes(self.intrinsic_dim)
60 }
61
62 fn bytes(&self) -> usize {
63 std::mem::size_of::<u8>() * self.nrows() * self.ncols()
64 }
65}
66
67impl<'a, const NBITS: usize> MatRef<'a, MinMaxMeta<NBITS>>
68where
69 Unsigned: Representation<NBITS>,
70{
71 pub(crate) fn as_canonical_view(&self) -> Result<rowmajor::Ref<'a, u8>, Overflow> {
73 let rows = self.num_vectors();
74 let dim = self.repr().intrinsic_dim();
75 let overflow = || Overflow::for_type::<u8>(rows, dim);
76 let stride = dim
77 .checked_mul(NBITS)
78 .ok_or_else(overflow)?
79 .div_ceil(8)
80 .checked_add(size_of::<MinMaxCompensation>())
81 .ok_or_else(overflow)?;
82 Overflow::check::<u8>(rows, stride)?;
83 let bytes = unsafe { std::slice::from_raw_parts(self.as_raw_ptr(), rows * stride) };
86 #[expect(
87 clippy::expect_used,
88 reason = "the checked dimensions match the slice length"
89 )]
90 Ok(rowmajor::Ref::try_from_data(bytes, rows, stride).expect("validated canonical rows"))
91 }
92}
93
94unsafe impl<const NBITS: usize> Repr for MinMaxMeta<NBITS>
99where
100 Unsigned: Representation<NBITS>,
101{
102 type Row<'a> = crate::minmax::DataRef<'a, NBITS>;
103
104 fn nrows(&self) -> usize {
105 self.nrows
106 }
107
108 fn layout(&self) -> Result<std::alloc::Layout, LayoutError> {
114 Ok(std::alloc::Layout::array::<u8>(
115 self.nrows() * self.ncols(),
116 )?)
117 }
118
119 unsafe fn get_row<'a>(self, ptr: NonNull<u8>, i: usize) -> Self::Row<'a> {
126 debug_assert!(i < self.nrows);
127 let len = self.ncols();
128
129 unsafe {
133 let row_ptr = ptr.as_ptr().add(i * len);
134 let slice = std::slice::from_raw_parts(row_ptr, len);
135
136 minmax::DataRef::<'a, NBITS>::from_canonical_unchecked(slice, self.intrinsic_dim)
137 }
138 }
139}
140
141unsafe impl<const NBITS: usize> ReprMut for MinMaxMeta<NBITS>
144where
145 Unsigned: Representation<NBITS>,
146{
147 type RowMut<'a> = crate::minmax::DataMutRef<'a, NBITS>;
148
149 unsafe fn get_row_mut<'a>(self, ptr: NonNull<u8>, i: usize) -> Self::RowMut<'a> {
157 debug_assert!(i < self.nrows);
158 let len = self.ncols();
159
160 unsafe {
165 let row_ptr = ptr.as_ptr().add(i * len);
166 let slice = std::slice::from_raw_parts_mut(row_ptr, len);
167
168 minmax::DataMutRef::<'a, NBITS>::from_canonical_front_mut_unchecked(
169 slice,
170 self.intrinsic_dim,
171 )
172 }
173 }
174}
175
176unsafe impl<const NBITS: usize> ReprOwned for MinMaxMeta<NBITS>
180where
181 Unsigned: Representation<NBITS>,
182{
183 unsafe fn drop(self, ptr: NonNull<u8>) {
188 let slice_ptr = std::ptr::slice_from_raw_parts_mut(ptr.as_ptr(), self.bytes());
189
190 let _ = unsafe { Box::from_raw(slice_ptr) };
193 }
194}
195
196unsafe impl<const NBITS: usize> NewOwned<Defaulted> for MinMaxMeta<NBITS>
200where
201 Unsigned: Representation<NBITS>,
202{
203 type Error = crate::error::Infallible;
204 fn new_owned(self, _: Defaulted) -> Result<Mat<Self>, Self::Error> {
205 let b: Box<[u8]> = (0..self.bytes()).map(|_| u8::default()).collect();
206
207 let ptr = utils::box_into_nonnull(b).cast::<u8>();
208
209 let mat = unsafe { Mat::from_raw_parts(self, ptr) };
212 Ok(mat)
213 }
214}
215
216unsafe impl<const NBITS: usize> NewRef<u8> for MinMaxMeta<NBITS>
219where
220 Unsigned: Representation<NBITS>,
221{
222 type Error = SliceError;
223
224 fn new_ref(self, slice: &[u8]) -> Result<MatRef<'_, Self>, Self::Error> {
225 let expected = self.bytes();
226 if slice.len() != expected {
227 return Err(SliceError::LengthMismatch {
228 expected,
229 found: slice.len(),
230 });
231 }
232
233 Ok(unsafe { MatRef::from_raw_parts(self, utils::as_nonnull(slice).cast::<u8>()) })
235 }
236}
237
238unsafe impl<const NBITS: usize> NewMut<u8> for MinMaxMeta<NBITS>
241where
242 Unsigned: Representation<NBITS>,
243{
244 type Error = SliceError;
245
246 fn new_mut(self, slice: &mut [u8]) -> Result<MatMut<'_, Self>, Self::Error> {
247 let expected = self.bytes();
248 if slice.len() != expected {
249 return Err(SliceError::LengthMismatch {
250 expected,
251 found: slice.len(),
252 });
253 }
254
255 Ok(unsafe { MatMut::from_raw_parts(self, utils::as_nonnull_mut(slice).cast::<u8>()) })
257 }
258}
259
260impl<'a, 'b, const NBITS: usize, T>
265 CompressInto<MatRef<'a, Standard<T>>, MatMut<'b, MinMaxMeta<NBITS>>> for MinMaxQuantizer
266where
267 T: Copy + Into<f32>,
268 Unsigned: Representation<NBITS>,
269{
270 type Error = InputContainsNaN;
271
272 type Output = ();
273
274 fn compress_into(
290 &self,
291 from: MatRef<'a, Standard<T>>,
292 mut to: MatMut<'b, MinMaxMeta<NBITS>>,
293 ) -> Result<(), Self::Error> {
294 assert_eq!(
295 from.num_vectors(),
296 to.num_vectors(),
297 "input and output must have the same number of vectors: {} != {}",
298 from.num_vectors(),
299 to.num_vectors()
300 );
301 assert_eq!(
302 from.vector_dim(),
303 self.dim(),
304 "input vectors must match quantizer dimension: {} != {}",
305 from.vector_dim(),
306 self.dim()
307 );
308 assert_eq!(
309 to.repr().intrinsic_dim(),
310 self.output_dim(),
311 "output intrinsic dimension must match quantizer output dimension: {} != {}",
312 to.repr().intrinsic_dim(),
313 self.output_dim()
314 );
315
316 for (from_row, to_row) in from.rows().zip(to.rows_mut()) {
317 let _ = <MinMaxQuantizer as CompressInto<&[T], DataMutRef<'_, NBITS>>>::compress_into(
319 self, from_row, to_row,
320 )?;
321 }
322
323 Ok(())
324 }
325}
326
327#[cfg(test)]
328mod tests {
329 use super::*;
330 use crate::algorithms::Transform;
331 use crate::algorithms::transforms::NullTransform;
332 use crate::minmax::vectors::DataRef;
333 use crate::num::Positive;
334 use diskann_utils::{Reborrow, ReborrowMut};
335 use std::num::NonZeroUsize;
336
337 const TEST_DIMS: &[usize] = &[1, 2, 3, 4, 7, 8, 16, 31, 32, 64];
339 const TEST_NVECS: &[usize] = &[1, 2, 3, 5, 10];
341
342 macro_rules! expand_to_bitrates {
344 ($name:ident, $func:ident) => {
345 #[test]
346 fn $name() {
347 $func::<1>();
348 $func::<2>();
349 $func::<4>();
350 $func::<8>();
351 }
352 };
353 }
354
355 fn make_quantizer(dim: usize) -> MinMaxQuantizer {
361 MinMaxQuantizer::new(
362 Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
363 Positive::new(1.0).unwrap(),
364 )
365 }
366
367 fn generate_test_data(num_vectors: usize, dim: usize) -> Vec<f32> {
369 (0..num_vectors * dim)
370 .map(|i| {
371 let row = i / dim;
372 let col = i % dim;
373 ((row as f32 + 1.0) * (col as f32 + 0.5)).sin() * 10.0 + (row as f32)
375 })
376 .collect()
377 }
378
379 fn compress_single_vector<const NBITS: usize>(
381 quantizer: &MinMaxQuantizer,
382 input: &[f32],
383 dim: usize,
384 ) -> Vec<u8>
385 where
386 Unsigned: Representation<NBITS>,
387 {
388 let row_bytes = Data::<NBITS>::canonical_bytes(dim);
389 let mut output = vec![0u8; row_bytes];
390 let output_ref =
391 crate::minmax::DataMutRef::<NBITS>::from_canonical_front_mut(&mut output, dim).unwrap();
392 let _ = <MinMaxQuantizer as CompressInto<&[f32], DataMutRef<'_, NBITS>>>::compress_into(
393 quantizer, input, output_ref,
394 )
395 .expect("single-vector compression should succeed");
396 output
397 }
398
399 mod construction {
404 use super::*;
405
406 fn test_new_owned<const NBITS: usize>()
408 where
409 Unsigned: Representation<NBITS>,
410 {
411 for &dim in TEST_DIMS {
412 for &num_vectors in TEST_NVECS {
413 let meta = MinMaxMeta::<NBITS>::new(num_vectors, dim);
414 let mat: Mat<MinMaxMeta<NBITS>> =
415 Mat::new(meta, Defaulted).expect("NewOwned should succeed");
416
417 assert_eq!(mat.num_vectors(), num_vectors);
419 assert_eq!(mat.repr().intrinsic_dim(), dim);
420
421 let expected_row_bytes = Data::<NBITS>::canonical_bytes(dim);
423 assert_eq!(mat.repr().nrows(), num_vectors);
424 assert_eq!(mat.repr().ncols(), expected_row_bytes);
425
426 let expected_bytes = expected_row_bytes * num_vectors;
428 let layout = mat.repr().layout().expect("layout should succeed");
429 assert_eq!(layout.size(), expected_bytes);
430
431 for i in 0..num_vectors {
433 let row = mat.get_row(i);
434 assert!(row.is_some(), "row {i} should exist");
435 assert_eq!(row.unwrap().len(), dim);
436 }
437
438 assert!(mat.get_row(num_vectors).is_none());
440 }
441 }
442 }
443
444 expand_to_bitrates!(new_owned, test_new_owned);
445
446 fn test_new_ref<const NBITS: usize>()
448 where
449 Unsigned: Representation<NBITS>,
450 {
451 for &dim in TEST_DIMS {
452 for &num_vectors in TEST_NVECS {
453 let meta = MinMaxMeta::<NBITS>::new(num_vectors, dim);
454 let expected_row_bytes = Data::<NBITS>::canonical_bytes(dim);
455 let expected_bytes = expected_row_bytes * num_vectors;
456
457 let data = vec![0u8; expected_bytes];
459 let mat_ref = MatRef::new(meta, &data);
460 assert!(mat_ref.is_ok(), "NewRef should succeed for correct size");
461 let mat_ref = mat_ref.unwrap();
462
463 assert_eq!(mat_ref.num_vectors(), num_vectors);
465 assert_eq!(mat_ref.repr().intrinsic_dim(), dim);
466
467 assert_eq!(mat_ref.repr().nrows(), num_vectors);
469 assert_eq!(mat_ref.repr().ncols(), expected_row_bytes);
470
471 let layout = mat_ref.repr().layout().expect("layout should succeed");
473 assert_eq!(layout.size(), expected_bytes);
474 }
475 }
476 }
477
478 expand_to_bitrates!(new_ref, test_new_ref);
479
480 fn test_new_mut<const NBITS: usize>()
482 where
483 Unsigned: Representation<NBITS>,
484 {
485 for &dim in TEST_DIMS {
486 for &num_vectors in TEST_NVECS {
487 let meta = MinMaxMeta::<NBITS>::new(num_vectors, dim);
488 let expected_row_bytes = Data::<NBITS>::canonical_bytes(dim);
489 let expected_bytes = expected_row_bytes * num_vectors;
490
491 let mut data = vec![0u8; expected_bytes];
492 let mat_mut = MatMut::new(meta, &mut data);
493 assert!(mat_mut.is_ok(), "NewMut should succeed for correct size");
494 let mat_mut = mat_mut.unwrap();
495
496 assert_eq!(mat_mut.num_vectors(), num_vectors);
498 assert_eq!(mat_mut.repr().intrinsic_dim(), dim);
499
500 assert_eq!(mat_mut.repr().nrows(), num_vectors);
502 assert_eq!(mat_mut.repr().ncols(), expected_row_bytes);
503
504 let layout = mat_mut.repr().layout().expect("layout should succeed");
506 assert_eq!(layout.size(), expected_bytes);
507 }
508 }
509 }
510
511 expand_to_bitrates!(new_mut, test_new_mut);
512
513 #[test]
514 fn slice_length_mismatch_errors() {
515 let dim = 4;
516 let num_vectors = 2;
517 let meta = MinMaxMeta::<8>::new(num_vectors, dim);
518 let expected_bytes = DataRef::<8>::canonical_bytes(dim) * num_vectors;
519
520 let short_data = vec![0u8; expected_bytes - 1];
522 let result = MatRef::new(meta, &short_data);
523 assert!(
524 matches!(result, Err(SliceError::LengthMismatch { .. })),
525 "should fail for too-short slice"
526 );
527
528 let long_data = vec![0u8; expected_bytes + 1];
530 let result = MatRef::new(meta, &long_data);
531 assert!(
532 matches!(result, Err(SliceError::LengthMismatch { .. })),
533 "should fail for too-long slice"
534 );
535
536 let mut short_mut = vec![0u8; expected_bytes - 1];
538 let result = MatMut::new(meta, &mut short_mut);
539 assert!(
540 matches!(result, Err(SliceError::LengthMismatch { .. })),
541 "MatMut should fail for too-short slice"
542 );
543 }
544 }
545
546 mod compress_into {
551 use super::*;
552
553 fn test_compress_matches_single<const NBITS: usize>()
556 where
557 Unsigned: Representation<NBITS>,
558 {
559 for &dim in TEST_DIMS {
560 for &num_vectors in TEST_NVECS {
561 let quantizer = make_quantizer(dim);
562 let input_data = generate_test_data(num_vectors, dim);
563
564 let input_view =
566 MatRef::new(Standard::new(num_vectors, dim).unwrap(), &input_data)
567 .expect("input view creation");
568
569 let mut multi_mat: Mat<MinMaxMeta<NBITS>> =
570 Mat::new(MinMaxMeta::new(num_vectors, dim), Defaulted)
571 .expect("output mat creation");
572
573 quantizer
574 .compress_into(input_view, multi_mat.reborrow_mut())
575 .expect("multi-vector compression");
576
577 for i in 0..num_vectors {
579 let row_input = &input_data[i * dim..(i + 1) * dim];
580 let expected_bytes =
581 compress_single_vector::<NBITS>(&quantizer, row_input, dim);
582
583 let actual_row = multi_mat.get_row(i).expect("row should exist");
584
585 let expected_ref = unsafe {
588 DataRef::<NBITS>::from_canonical_unchecked(&expected_bytes, dim)
589 };
590 assert_eq!(
591 actual_row.meta(),
592 expected_ref.meta(),
593 "metadata mismatch at row {i}, dim={dim}, num_vectors={num_vectors}, NBITS={NBITS}"
594 );
595
596 for j in 0..dim {
598 assert_eq!(
599 actual_row.vector().get(j).unwrap(),
600 expected_ref.vector().get(j).unwrap(),
601 "quantized value mismatch at row {i}, col {j}"
602 );
603 }
604 }
605 }
606 }
607 }
608
609 expand_to_bitrates!(compress_matches_single, test_compress_matches_single);
610
611 fn test_row_iteration<const NBITS: usize>()
613 where
614 Unsigned: Representation<NBITS>,
615 {
616 let dim = 8;
617 let num_vectors = 5;
618 let quantizer = make_quantizer(dim);
619 let input_data = generate_test_data(num_vectors, dim);
620
621 let input_view = MatRef::new(Standard::new(num_vectors, dim).unwrap(), &input_data)
622 .expect("input view");
623
624 let mut mat: Mat<MinMaxMeta<NBITS>> =
625 Mat::new(MinMaxMeta::new(num_vectors, dim), Defaulted).expect("mat creation");
626
627 quantizer
628 .compress_into(input_view, mat.reborrow_mut())
629 .expect("compression");
630
631 let view = mat.reborrow();
633 let mut count = 0;
634 for row in view.rows() {
635 assert_eq!(row.len(), dim, "row should have correct dimension");
636 count += 1;
637 }
638 assert_eq!(count, num_vectors);
639 }
640
641 expand_to_bitrates!(row_iteration, test_row_iteration);
642 }
643
644 mod error_cases {
649 use super::*;
650
651 #[test]
652 #[should_panic(expected = "input and output must have the same number of vectors")]
653 fn compress_into_vector_count_mismatch() {
654 const NBITS: usize = 8;
655 let dim = 4;
656 let quantizer = make_quantizer(dim);
657
658 let input_data = generate_test_data(3, dim);
660 let input_view =
661 MatRef::new(Standard::new(3, dim).unwrap(), &input_data).expect("input view");
662
663 let mut mat: Mat<MinMaxMeta<NBITS>> =
665 Mat::new(MinMaxMeta::new(2, dim), Defaulted).expect("mat creation");
666
667 let _ = quantizer.compress_into(input_view, mat.reborrow_mut());
668 }
669
670 #[test]
671 #[should_panic(expected = "input vectors must match quantizer dimension")]
672 fn compress_into_input_dim_mismatch() {
673 const NBITS: usize = 8;
674 let quantizer = make_quantizer(4); let input_data = generate_test_data(2, 8);
678 let input_view =
679 MatRef::new(Standard::new(2, 8).unwrap(), &input_data).expect("input view");
680
681 let mut mat: Mat<MinMaxMeta<NBITS>> =
683 Mat::new(MinMaxMeta::new(2, 4), Defaulted).expect("mat creation");
684
685 let _ = quantizer.compress_into(input_view, mat.reborrow_mut());
686 }
687
688 #[test]
689 #[should_panic(
690 expected = "output intrinsic dimension must match quantizer output dimension"
691 )]
692 fn compress_into_output_dim_mismatch() {
693 const NBITS: usize = 8;
694 let quantizer = make_quantizer(4);
695
696 let input_data = generate_test_data(2, 4);
698 let input_view =
699 MatRef::new(Standard::new(2, 4).unwrap(), &input_data).expect("input view");
700
701 let row_bytes = Data::<NBITS>::canonical_bytes(8);
703 let mut output_data = vec![0u8; 2 * row_bytes];
704 let output_view =
705 MatMut::new(MinMaxMeta::<NBITS>::new(2, 8), &mut output_data).expect("output view");
706
707 let _ = quantizer.compress_into(input_view, output_view);
708 }
709 }
710}