1use yo_common::{Addr, Code, Error, Result};
51use yo_format::image::{
52 Chain, ImageHeader, Keys, PostingHeader, get_floats, image_kind, image_len, key_entry_len,
53 metric, posting_len, put_floats, put_key, put_partition,
54};
55use yo_format::{get_f32, get_u64, put_u64};
56use yo_kv::cold::{self, Blocks, Scratch};
57use yo_shape::Metric;
58
59use crate::collection::{Collection, check_metric};
60use crate::partition::{Partitions, Tuning};
61use crate::rabitq::{Bits, Coded};
62
63pub trait Stored {
70 fn get(&self, key: &[u8], into: &mut [f32]) -> bool;
73}
74
75#[derive(Debug)]
77pub struct Restored {
78 pub collection: Collection,
80 pub missing: usize,
88}
89
90impl Collection {
91 pub fn save<B: Blocks>(&self, blocks: &mut B, scratch: &mut Scratch) -> Result<Chain> {
103 let index = self.index();
104 let dim = index.dim();
105 let width = index.quantizer().code_bytes();
106 let count = index.partitions();
107
108 let mut buf = Vec::new();
109 let mut root = vec![0u8; image_len(as_u32(count)?)?];
110
111 for p in 0..count {
114 let (ids, tags, codes, meta, stuck) = index.posting_parts(p);
115 let head = PostingHeader {
116 count: as_u32(ids.len())?,
117 code_bytes: as_u32(width)?,
118 stuck: as_u32(stuck)?,
119 };
120 buf.clear();
121 buf.resize(posting_len(head.count, head.code_bytes)?, 0);
122 head.encode(&mut buf)?;
123 for (i, &id) in ids.iter().enumerate() {
124 put_u64(&mut buf, head.ids_at() + i * 8, id);
125 }
126 for (i, &tag) in tags.iter().enumerate() {
127 put_u64(&mut buf, head.tags_at() + i * 8, tag);
128 }
129 let at = head.codes_at();
130 buf[at..at + codes.len()].copy_from_slice(codes);
131 for (i, m) in meta.iter().enumerate() {
132 let at = head.meta_at() + i * 16;
133 put_floats(&mut buf[at..], &[m.norm, m.scale, m.lo, m.delta])?;
134 }
135 put_partition(&mut root, as_u32(p)?, write(blocks, &buf, scratch)?)?;
136 }
137
138 buf.clear();
141 buf.resize(index.all_centroids().len() * 4, 0);
142 put_floats(&mut buf, index.all_centroids())?;
143 let centroids = write(blocks, &buf, scratch)?;
144
145 buf.clear();
146 for (key, &id) in self.id_table().iter() {
147 let at = buf.len();
148 buf.resize(at + key_entry_len(key.len())?, 0);
149 put_key(&mut buf[at..], id, key)?;
150 }
151 let keys = write(blocks, &buf, scratch)?;
152
153 let tuning = index.tuning();
154 let head = ImageHeader {
155 kind: image_kind::VECTOR,
156 bits: as_u32(index.quantizer().bits().count())? as u8,
157 metric: metric_byte(self.metric()),
158 dim: as_u32(dim)?,
159 partitions: as_u32(count)?,
160 seed: index.quantizer().seed(),
161 members: self.len() as u64,
162 slots: as_u32(self.slots())?,
163 posting: as_u32(tuning.posting)?,
164 probe: as_u32(tuning.probe)?,
165 rerank: as_u32(tuning.rerank)?,
166 sweep: as_u32(tuning.sweep)?,
167 widen: as_u32(tuning.widen)?,
168 spill: as_u32(tuning.spill)?,
169 slack: tuning.slack,
170 patience: as_u32(tuning.patience)?,
171 centroids,
172 keys,
173 };
174 head.encode(&mut root)?;
175 write(blocks, &root, scratch)
176 }
177
178 pub fn load<B: Blocks>(blocks: &mut B, at: Chain, stored: &impl Stored) -> Result<Restored> {
187 let mut buf = Vec::new();
188 read(blocks, at, &mut buf)?;
189 let head = ImageHeader::decode(&buf)?;
190 if head.kind != image_kind::VECTOR {
191 return Err(
192 Error::new(Code::Corrupt, "that image is not a vector index")
193 .with_detail(format!("kind={}", head.kind)),
194 );
195 }
196 let bits = match head.bits {
197 1 => Bits::One,
198 _ => Bits::Four,
199 };
200 let metric = metric_of(head.metric)?;
201 check_metric(metric)?;
202 let dim = head.dim as usize;
203 let root = std::mem::take(&mut buf);
204
205 let mut index = Partitions::new(
206 dim,
207 bits,
208 head.seed,
209 Tuning {
210 posting: head.posting as usize,
211 probe: head.probe as usize,
212 rerank: head.rerank as usize,
213 sweep: head.sweep as usize,
214 widen: head.widen as usize,
215 spill: head.spill as usize,
216 slack: head.slack,
217 patience: head.patience as usize,
218 },
219 );
220 let width = index.quantizer().code_bytes();
221
222 blocks.release();
223 read(blocks, head.centroids, &mut buf)?;
224 let mut centroids = vec![0f32; head.partitions as usize * dim];
225 get_floats(&buf, &mut centroids)?;
226
227 let mut members = Vec::new();
228 for p in 0..head.partitions {
229 blocks.release();
230 read(blocks, yo_format::image::get_partition(&root, p)?, &mut buf)?;
231 let post = PostingHeader::decode(&buf)?;
232 if post.code_bytes as usize != width {
233 return Err(
234 Error::new(Code::Corrupt, "a partition's codes are the wrong width")
235 .with_detail(format!("code_bytes={} want={width}", post.code_bytes)),
236 );
237 }
238 let n = post.count as usize;
239 let mut ids = Vec::with_capacity(n);
240 let mut tags = Vec::with_capacity(n);
241 let mut meta = Vec::with_capacity(n);
242 for i in 0..n {
243 let id = get_u64(&buf, post.ids_at() + i * 8);
244 if id >= u64::from(head.slots) {
245 return Err(Error::new(Code::Corrupt, "a member's id is past the table")
246 .with_detail(format!("id={id} slots={}", head.slots)));
247 }
248 ids.push(id);
249 tags.push(get_u64(&buf, post.tags_at() + i * 8));
250 let at = post.meta_at() + i * 16;
251 meta.push(Coded {
252 norm: get_f32(&buf, at),
253 scale: get_f32(&buf, at + 4),
254 lo: get_f32(&buf, at + 8),
255 delta: get_f32(&buf, at + 12),
256 });
257 }
258 let codes = buf[post.codes_at()..post.meta_at()].to_vec();
259 let at = p as usize * dim;
260 index.absorb(
261 ¢roids[at..at + dim],
262 ids,
263 tags,
264 codes,
265 meta,
266 post.stuck as usize,
267 )?;
268 members.push(n);
269 }
270 index.finish_image();
271
272 blocks.release();
273 read(blocks, head.keys, &mut buf)?;
274 let mut collection = Collection::from_image(index, metric, head.slots as usize);
275 let mut vector = vec![0f32; dim];
276 let mut missing = 0;
277 let mut named = 0u64;
278 let mut walk = Keys::new(&buf);
279 for (id, key) in walk.by_ref() {
280 named += 1;
281 if !collection.holds(id) {
282 return Err(Error::new(
283 Code::Corrupt,
284 "the key table names an id no partition has",
285 )
286 .with_detail(format!("id={id}")));
287 }
288 if stored.get(key, &mut vector) {
289 collection.restore(key, id, &vector)?;
290 } else {
291 collection.forget(id);
292 missing += 1;
293 }
294 }
295 if !walk.done() {
296 return Err(Error::new(Code::Corrupt, "the key table ends mid entry"));
297 }
298 if named != head.members {
299 return Err(Error::new(
300 Code::Corrupt,
301 "the key table is not the length the header says",
302 )
303 .with_detail(format!("keys={named} members={}", head.members)));
304 }
305 collection.seal();
306 Ok(Restored {
307 collection,
308 missing,
309 })
310 }
311}
312
313fn write<B: Blocks>(blocks: &mut B, bytes: &[u8], scratch: &mut Scratch) -> Result<Chain> {
315 let chain = cold::write(blocks, bytes, scratch)?;
316 Ok(Chain {
317 at: chain.at.to_bits(),
318 len: chain.len,
319 })
320}
321
322fn read<B: Blocks>(blocks: &B, at: Chain, out: &mut Vec<u8>) -> Result<()> {
324 out.clear();
325 let reader = cold::Reader::open(
326 blocks,
327 cold::Chain {
328 at: Addr::from_bits(at.at),
329 len: at.len,
330 },
331 )?;
332 for piece in reader.range(0, at.len) {
333 out.extend_from_slice(piece?);
334 }
335 Ok(())
336}
337
338fn metric_byte(m: Metric) -> u8 {
340 match m {
341 Metric::L2 => metric::L2,
342 Metric::Cosine => metric::COSINE,
343 Metric::Ip => metric::IP,
344 Metric::Hamming => metric::HAMMING,
345 }
346}
347
348fn metric_of(b: u8) -> Result<Metric> {
350 match b {
351 metric::L2 => Ok(Metric::L2),
352 metric::COSINE => Ok(Metric::Cosine),
353 metric::IP => Ok(Metric::Ip),
354 metric::HAMMING => Ok(Metric::Hamming),
355 _ => Err(Error::new(Code::Corrupt, "unknown metric in an image")
356 .with_detail(format!("metric={b}"))),
357 }
358}
359
360fn as_u32(n: usize) -> Result<u32> {
362 u32::try_from(n).map_err(|_| {
363 Error::new(Code::Full, "that count does not fit in an image").with_detail(format!("n={n}"))
364 })
365}
366
367#[cfg(test)]
368mod tests {
369 use super::*;
370 use std::collections::HashMap;
371
372 struct Mem {
375 blobs: Vec<Vec<u8>>,
376 }
377
378 impl Mem {
379 fn new() -> Mem {
380 Mem { blobs: Vec::new() }
381 }
382 }
383
384 impl Blocks for Mem {
385 fn put(&mut self, bytes: &[u8]) -> Result<Addr> {
386 self.blobs.push(bytes.to_vec());
387 Ok(Addr::new(
388 yo_common::Space::Log,
389 (self.blobs.len() - 1) as u64,
390 ))
391 }
392
393 fn get(&self, at: Addr) -> Result<&[u8]> {
394 self.blobs
395 .get(at.offset() as usize)
396 .map(Vec::as_slice)
397 .ok_or_else(|| Error::new(Code::NotFound, "no such block"))
398 }
399
400 fn bytes(&self) -> u64 {
401 self.blobs.iter().map(|b| b.len() as u64).sum()
402 }
403 }
404
405 struct Table(HashMap<Vec<u8>, Vec<f32>>);
407
408 impl Table {
409 fn of(c: &Collection) -> Table {
410 let mut m = HashMap::new();
411 for key in c.keys() {
412 m.insert(
413 key.to_vec(),
414 c.get(key).expect("a key it just named").to_vec(),
415 );
416 }
417 Table(m)
418 }
419
420 fn without(mut self, key: &[u8]) -> Table {
421 self.0.remove(key);
422 self
423 }
424 }
425
426 impl Stored for Table {
427 fn get(&self, key: &[u8], into: &mut [f32]) -> bool {
428 let Some(v) = self.0.get(key) else {
429 return false;
430 };
431 into.copy_from_slice(v);
432 true
433 }
434 }
435
436 fn corpus(dim: usize, n: usize, seed: u64) -> Vec<(Vec<u8>, Vec<f32>)> {
439 let mut state = seed | 1;
440 let mut next = move || {
441 state ^= state << 13;
442 state ^= state >> 7;
443 state ^= state << 17;
444 (state >> 11) as f32 / (1u64 << 53) as f32 - 0.5
445 };
446 (0..n)
447 .map(|i| {
448 let key = format!("k{i}").into_bytes();
449 let v: Vec<f32> = (0..dim).map(|_| next()).collect();
450 (key, v)
451 })
452 .collect()
453 }
454
455 fn clustered(dim: usize, n: usize, clusters: usize, seed: u64) -> Vec<(Vec<u8>, Vec<f32>)> {
461 let centres = corpus(dim, clusters, seed);
462 corpus(dim, n, seed ^ 0x9e37)
463 .into_iter()
464 .enumerate()
465 .map(|(i, (key, off))| {
466 let centre = ¢res[i % clusters].1;
467 let v = centre.iter().zip(&off).map(|(c, o)| c + o * 0.6).collect();
468 (key, v)
469 })
470 .collect()
471 }
472
473 fn built(dim: usize, n: usize, metric: Metric) -> Collection {
474 built_from(dim, metric, Tuning::default(), corpus(dim, n, 42))
475 }
476
477 fn built_from(
478 dim: usize,
479 metric: Metric,
480 tuning: Tuning,
481 vectors: Vec<(Vec<u8>, Vec<f32>)>,
482 ) -> Collection {
483 let mut c = Collection::new(dim, metric).expect("a collection");
484 c.retune(tuning);
485 for (i, (key, v)) in vectors.into_iter().enumerate() {
486 c.put_tagged(&key, &v, 1 << (i % 8)).expect("put");
487 }
488 c
489 }
490
491 fn round_trip(c: &Collection, stored: &impl Stored) -> Restored {
492 let mut mem = Mem::new();
493 let mut scratch = Scratch::new();
494 let at = c.save(&mut mem, &mut scratch).expect("saved");
495 Collection::load(&mut mem, at, stored).expect("loaded")
496 }
497
498 #[test]
499 fn a_collection_comes_back_answering_the_same_questions() {
500 let c = built(32, 900, Metric::L2);
501 let back = round_trip(&c, &Table::of(&c)).collection;
502
503 assert_eq!(back.len(), c.len());
504 assert_eq!(back.dim(), c.dim());
505 assert_eq!(back.metric(), c.metric());
506 assert!(
507 c.partitions() > 1,
508 "a collection that never split proves nothing"
509 );
510 assert_eq!(
511 back.partitions(),
512 c.partitions(),
513 "the shape is the thing an image exists to keep"
514 );
515 assert_eq!(back.tuning(), c.tuning());
516
517 for (_, q) in corpus(32, 50, 7) {
522 assert_eq!(
523 back.search(&q, 10, None).expect("search"),
524 c.search(&q, 10, None).expect("search"),
525 "a query came back differently after a round trip"
526 );
527 assert_eq!(
528 back.search_where(&q, 10, None, &crate::Signature::from_bits(1 << 3))
529 .expect("search"),
530 c.search_where(&q, 10, None, &crate::Signature::from_bits(1 << 3))
531 .expect("search"),
532 "a filtered query came back differently, so a tag moved"
533 );
534 }
535 }
536
537 #[test]
546 fn the_boundary_copies_survive_a_round_trip() {
547 let tuning = Tuning {
548 slack: 0.25,
549 ..Tuning::default()
550 };
551 let c = built_from(32, Metric::L2, tuning, clustered(32, 3000, 12, 42));
552 assert!(
553 c.entries() > c.len(),
554 "a collection with no copies in it proves nothing here"
555 );
556 let back = round_trip(&c, &Table::of(&c)).collection;
557 assert_eq!(back.entries(), c.entries(), "a copy was lost");
558 assert_eq!(back.len(), c.len(), "a copy came back as a member");
559 assert_eq!(back.partitions(), c.partitions());
560 }
561
562 #[test]
563 fn every_vector_and_every_tag_survives() {
564 let c = built(16, 400, Metric::Cosine);
565 let back = round_trip(&c, &Table::of(&c)).collection;
566 for key in c.keys() {
567 assert_eq!(
568 back.get(key).map(<[f32]>::to_vec),
569 c.get(key).map(<[f32]>::to_vec),
570 "the vector under a key changed, bit for bit"
571 );
572 assert_eq!(back.tag(key), c.tag(key), "a tag was lost");
573 }
574 }
575
576 #[test]
577 fn a_reloaded_collection_takes_writes_where_it_left_off() {
578 let mut c = built(16, 300, Metric::L2);
579 let mut back = round_trip(&c, &Table::of(&c)).collection;
580
581 assert!(c.remove(b"k7"));
584 assert!(back.remove(b"k7"));
585 for (key, v) in corpus(16, 40, 99) {
586 let key = [b"new-".as_slice(), &key].concat();
587 c.put(&key, &v).expect("put");
588 back.put(&key, &v).expect("put");
589 }
590 assert_eq!(back.len(), c.len());
591 for key in c.keys() {
592 assert!(back.contains(key), "a key written after a load is missing");
593 }
594 let q = c.get(b"k1").expect("a vector").to_vec();
595 assert_eq!(
596 back.search(&q, 5, None).expect("search"),
597 c.search(&q, 5, None).expect("search")
598 );
599 }
600
601 #[test]
607 fn a_partition_that_gave_up_splitting_does_not_try_again_after_a_load() {
608 let mut c = Collection::new(8, Metric::L2).expect("a collection");
609 for i in 0..1000 {
610 c.put(format!("same{i}").as_bytes(), &[0.5; 8])
611 .expect("put");
612 }
613 assert_eq!(c.maintain(1 << 20), 0, "it has already given up");
614
615 let mut back = round_trip(&c, &Table::of(&c)).collection;
616 assert_eq!(
617 back.maintain(1 << 20),
618 0,
619 "the load forgot that the split was hopeless and went looking again"
620 );
621 assert_eq!(back.partitions(), c.partitions());
622 }
623
624 #[test]
625 fn an_empty_collection_is_an_image_too() {
626 let c = Collection::new(8, Metric::L2).expect("a collection");
627 let back = round_trip(&c, &Table::of(&c)).collection;
628 assert!(back.is_empty());
629 assert_eq!(back.partitions(), 0);
630 assert!(back.search(&[0.0; 8], 4, None).expect("search").is_empty());
631 }
632
633 #[test]
638 fn a_section_longer_than_a_chunk_is_still_one_section() {
639 let mut c = Collection::new(8, Metric::L2).expect("a collection");
640 for (i, (_, v)) in corpus(8, 5000, 11).into_iter().enumerate() {
641 c.put(format!("key{i}").as_bytes(), &v).expect("put");
642 }
643
644 let mut mem = Mem::new();
645 let mut scratch = Scratch::new();
646 let at = c.save(&mut mem, &mut scratch).expect("saved");
647 assert!(
648 mem.blobs.len() > c.partitions() + 3,
649 "no section was cut up, so this proves nothing about chains"
650 );
651
652 let back = Collection::load(&mut mem, at, &Table::of(&c))
653 .expect("loaded")
654 .collection;
655 assert_eq!(back.len(), c.len());
656 for key in c.keys() {
657 assert!(
658 back.contains(key),
659 "a key on the far side of a chunk is gone"
660 );
661 }
662 }
663
664 #[test]
665 fn a_key_the_store_cannot_produce_is_dropped_and_counted() {
666 let c = built(16, 200, Metric::L2);
667 let restored = round_trip(&c, &Table::of(&c).without(b"k5"));
668 assert_eq!(restored.missing, 1);
669 let back = restored.collection;
670 assert_eq!(back.len(), c.len() - 1);
671 assert!(!back.contains(b"k5"));
672 let mut back = back;
675 back.put(b"k5", &[1.0; 16]).expect("put");
676 assert_eq!(back.len(), c.len());
677 assert_eq!(back.get(b"k5"), Some([1.0f32; 16].as_slice()));
678 }
679
680 #[test]
681 fn an_image_that_says_something_impossible_is_refused() {
682 let c = built(8, 120, Metric::L2);
683 let mut mem = Mem::new();
684 let mut scratch = Scratch::new();
685 let at = c.save(&mut mem, &mut scratch).expect("saved");
686 assert!(Collection::load(&mut mem, at, &Table::of(&c)).is_ok());
687
688 let root = mem.blobs.len() - 1;
691 for (at_byte, to) in [(4usize, 9u8), (5, 2), (6, 9), (7, 1)] {
692 let mut broken = Mem {
693 blobs: mem.blobs.clone(),
694 };
695 broken.blobs[root][at_byte] = to;
696 assert!(
697 Collection::load(&mut broken, at, &Table::of(&c)).is_err(),
698 "byte {at_byte} of the root was believed"
699 );
700 }
701
702 let mut lying = Mem {
705 blobs: mem.blobs.clone(),
706 };
707 put_u64(&mut lying.blobs[root], 24, 3);
708 assert!(Collection::load(&mut lying, at, &Table::of(&c)).is_err());
709 }
710}