yo_vector/partition.rs
1//! The partition index, and the in place update protocol that means it never
2//! has to be rebuilt (`10` sections 2, 4 and 5).
3//!
4//! A vector index is two decisions and the quantiser was the easy one. This is
5//! the other: what holds the codes, and what happens to it when the collection
6//! changes.
7//!
8//! The answer is not a graph. Redis shipped HNSW vector sets in 8.0 in May 2025
9//! and they were still beta three minor releases later, which is the vendor's
10//! own evidence about how that goes. The reason is the update path: a graph
11//! index tombstones a delete, degrades as the tombstones pile up, and only gets
12//! better again when somebody rebuilds it, which on a collection anyone cares
13//! about is an outage with a nicer name.
14//!
15//! So this is partitions. Every vector belongs to the partition whose centroid
16//! it is nearest, the centroids are resident, and a partition's members are a
17//! flat run of codes that the scan walks end to end. An insert is an append. A
18//! delete takes a member out and moves the last one into the hole. Neither one
19//! touches anything else.
20//!
21//! # Search
22//!
23//! Rank the centroids, take the nearest `probe` of them, scan those partitions
24//! with the estimator, keep the best `rerank` candidates, and then look up the
25//! full precision vectors for exactly those and measure them properly. Rerank
26//! costs nothing structurally here, because the vector is already in the record
27//! log at an address the id resolves to. Every other system that quantises has
28//! to keep the raw vectors somewhere on purpose.
29//!
30//! The scan is the shape hardware likes, a linear walk over contiguous bytes,
31//! and that is why an index of this family beats a graph on a modern core even
32//! though it looks at more candidates.
33//!
34//! # It never rebuilds
35//!
36//! This is SPFresh's LIRE, and it is four bounded operations rather than a
37//! background rebuild.
38//!
39//! A posting that grows past twice its target splits, by two means over its own
40//! members. A posting that falls under a quarter of its target merges, by
41//! handing its members to whichever centroid is nearest now. Both are bounded
42//! work on one partition.
43//!
44//! The third is the one that matters and it is what LIRE actually contributes.
45//! After a split, the members of the partitions *around* the one that split may
46//! now be nearer one of the two new centroids than the one they are filed
47//! under. Nobody told them, and a plain partitioned index just lets that drift,
48//! which is why a plain partitioned index measures beautifully on a freshly
49//! built corpus and badly after a week of writes. So a split is followed by a
50//! sweep of the neighbouring partitions, and anything whose nearest centroid
51//! has changed is moved. The test for this writes a stream and checks recall at
52//! the end of it rather than on a fresh build, because a fresh build is exactly
53//! the measurement that hides the problem.
54//!
55//! # Everything here is in rotated space
56//!
57//! The rotation is linear, so `rotate(v - c)` is `rotate(v) - rotate(c)`, and
58//! distances and angles come through it unchanged. That means a centroid can be
59//! stored already rotated and a query can be rotated once, and then meeting a
60//! partition is a subtraction rather than another rotation. The rotation is the
61//! expensive half of preparing a query, so on a search that probes tens of
62//! partitions this is most of what preparation costs.
63//!
64//! # What it costs to build
65//!
66//! `examples/ingest.rs` measures the rate at every doubling and splits it
67//! between the insert and the maintenance, because a rate that falls as the
68//! collection grows and a rate that is just low need different work and a single
69//! number cannot tell them apart. On 128 dimensional vectors on one core of a
70//! 13th Gen Intel Core i9-13900K with nothing else running:
71//!
72//! ```text
73//! at partitions a second insert maintain touched
74//! 12500 36 132060 20.7% 79.3% 5.3
75//! 50000 132 94166 31.0% 69.0% 6.9
76//! 200000 595 67764 49.1% 50.9% 5.5
77//! 800000 2141 64107 52.8% 47.2% 4.4
78//! 1600000 4337 59040 53.7% 46.3% 4.4
79//! ```
80//!
81//! `touched` is how many vectors maintenance moved or looked at per vector
82//! inserted. It is flat, and that is the number which says the update protocol
83//! is doing bounded work rather than quietly turning into a rebuild.
84//!
85//! Five fixes got it there and they were five different problems.
86//! Maintenance was 80 percent of the time and most of it was `sweep` measuring
87//! every member it looked at against every centroid in the collection, which is
88//! not what LIRE says and is several full scans per vector inserted. The insert
89//! was the other half and it was a scan over every centroid by definition, which
90//! is why the coarse layer in `src/coarse.rs` is there, and that file is where
91//! the reasoning about it lives. Before either fix, the rate halved on every
92//! doubling and was 13563 a second by 800 thousand.
93//!
94//! The third was the squared distance itself, which by then was half of an
95//! entire ingest across the two copies of it that existed, and it was slow for
96//! a reason that had nothing to do with the index: a bounds check the compiler
97//! could not remove was stopping the loop vectorising. There is now one copy in
98//! `src/dist.rs` and that file is where the reasoning lives. It doubled the
99//! rate on its own, 21183 a second to 40318 over the whole 1.6 million on the
100//! same machine, and it cut the insert half by three times, 45.9 seconds to
101//! 15.3, which is why insert and maintenance have swapped places in that table.
102//!
103//! The fourth was `Partitions::job`, which used to ask its two questions by
104//! walking every partition twice, once per vector inserted. That is the same
105//! quadratic the coarse layer exists to remove, hiding one level up, and by 1.6
106//! million vectors it was around a fifth of an ingest spent deciding there was
107//! nothing to do. It is the two candidate lists now.
108//!
109//! The fifth was the rotation, in `src/rotate.rs`, which unpacked a sign bit
110//! with a shift and a mask inside the loop and then branched on a random bit
111//! per pair. Turning a pair is the same as flipping the sign of its second
112//! coordinate, so the branch folds into the sign table at build time. Those two
113//! together took the whole 1.6 million from 40318 a second to 64647, and the
114//! maintenance half from 24.4 seconds to 11.6.
115//!
116//! So G13's fifty thousand a second per core is met on the machine it is called
117//! on, and it is met at every size in the table rather than only at the small
118//! end. The two halves are close to even now, 53.7 percent insert against 46.3
119//! percent maintenance at the far end, so neither one is the obvious next thing
120//! to go and look at.
121//!
122//! # What is not here yet
123//!
124//! A checkpoint that writes the image out. [`crate::image`] is the layout and
125//! the two halves of the round trip, and the seam it comes back through is a
126//! pair of crate private calls further down this file, so an index survives a
127//! restart without requantising anything. What is still missing is the shard
128//! side: deciding when to write one, and pointing a checkpoint entry at it.
129//!
130//! MS-MARCO-v2. SIFT1M on a 13900K now gets recall 0.9597 at probe 64 rerank
131//! 16 with p50 at 638 us and p99 at 776 us, so both halves of G12 are met on
132//! that dataset, and the same run before the `src/dist.rs` change was 808 us
133//! and 996 us for the same recall, which was inside the millisecond by so
134//! little that nobody should have called it. The other dataset the gate names
135//! has not been run.
136//!
137//! `examples/search.rs` is the breakdown of where the remaining time goes and
138//! the answer is that two thirds of it is the estimator meeting one code at a
139//! time.
140//!
141//! The commands that put all of this on the wire are the rest of M6.
142
143use std::collections::{HashMap, HashSet};
144
145use yo_common::{Code, Error, Result};
146
147use crate::coarse::Coarse;
148use crate::dist::sqdist;
149use crate::rabitq::{Bits, Coded, Quantizer};
150
151/// Where the full precision vectors live.
152///
153/// A real collection answers this out of the record log, which already holds
154/// the vector at an address the id resolves to. A test answers it out of a map.
155/// Either way the index itself never stores a raw vector, which is the whole
156/// point of quantising one.
157pub trait Vectors {
158 /// Write the vector `id` stands for into `into` and say so, or say that the
159 /// id is gone.
160 ///
161 /// An id that is gone is dropped from the index the next time maintenance
162 /// walks over it, so a collection that deletes from the log without telling
163 /// the index heals rather than lying.
164 fn get(&self, id: u64, into: &mut [f32]) -> bool;
165}
166
167/// The knobs, all of which have a defensible default and none of which anybody
168/// should have to touch.
169#[derive(Debug, Clone, Copy, PartialEq)]
170pub struct Tuning {
171 /// How many members a partition wants. It splits past twice this and merges
172 /// under a quarter of it.
173 ///
174 /// This is what sets how many partitions a collection ends up with, and so
175 /// it trades the cost of ranking centroids against the cost of scanning a
176 /// posting. A few hundred is where those two are near enough even.
177 pub posting: usize,
178 /// How many partitions a search scans.
179 pub probe: usize,
180 /// How many candidates are reranked per answer asked for.
181 ///
182 /// Four is the number the recall table was measured at: one bit codes put
183 /// the true ten inside the best forty better than 98 times in a hundred.
184 pub rerank: usize,
185 /// How many neighbouring partitions a split sweeps for members that should
186 /// move.
187 ///
188 /// This is the cost of never drifting. Zero would make a split free and
189 /// would make recall fall off over a long write stream, which is the thing
190 /// this index exists to not do.
191 pub sweep: usize,
192 /// How much further than `probe` a filtered search will go looking when the
193 /// filter is selective enough that the nearest partitions do not hold `k`
194 /// members that pass, as a multiple of `probe`.
195 ///
196 /// This is the only knob here with a genuinely hard trade behind it. Too
197 /// small and a filter matching one document in a thousand returns nothing
198 /// while the answer sat two partitions further out. Too large and the same
199 /// filter reads the whole collection to prove there is nothing there.
200 pub widen: usize,
201 /// How many partitions in a row may add nothing to the answer before the
202 /// search stops reading, once it has enough candidates to answer with.
203 ///
204 /// [`Tuning::probe`] is a budget every query spends whether it needs to or
205 /// not, and queries do not need the same amount. A query sitting deep inside
206 /// one partition has found everything it is going to find after two or three
207 /// of them, and a query on a boundary is still turning up better answers
208 /// forty partitions in. This is what lets one search cost what it needs
209 /// rather than what the slowest query needs, and it is the whole of the
210 /// difference between a mean probe depth and a fixed one.
211 ///
212 /// It keys off the answer rather than off the geometry, which is deliberate.
213 /// The obvious rule is to stop once the next centroid is more than some
214 /// fraction further away than the nearest one, and that rule is useless
215 /// here: the measurement is on the private `spill_into`, which is where the
216 /// same rule was tried and dropped on the write path, and the short version
217 /// is that distances concentrate, every centroid a query can see is
218 /// within a few percent of every other, and there is no setting of the
219 /// fraction between pruning nothing and pruning everything.
220 ///
221 /// A partition counts as adding nothing when not one of its members was good
222 /// enough to displace an answer already held. The count resets the moment one
223 /// is, so a run of empty partitions followed by a good one buys the search
224 /// its patience back. Zero switches this off and every search reads `probe`
225 /// partitions.
226 ///
227 /// It cannot change how many candidates come back, only which ones, because
228 /// it is only allowed to fire once there are already enough. A filtered
229 /// search that is widening because it does not have enough is never cut off
230 /// by it.
231 ///
232 /// # Where the default comes from
233 ///
234 /// Eight, which is the same as the default `probe`, and that is not a
235 /// coincidence: a search that only reads eight partitions cannot have eight
236 /// quiet ones in a row before it runs out, so the default settings are the
237 /// settings this does nothing under. It starts to matter exactly when
238 /// somebody raises `probe`, which is when it should.
239 ///
240 /// On SIFT1M at `probe` 128 and `rerank` 16, where the fixed sweep recalls
241 /// 0.9757 reading all 128, eight reads 96.9 of them for 0.9750, four reads
242 /// 65.6 for 0.9704, and three reads 53.3 for 0.9641, against a fixed `probe`
243 /// of 64 which reads all 64 for 0.9665. So a quarter of the reads go away for
244 /// seven ten thousandths of recall, and the settings in between fill in a
245 /// ladder that `probe` can only climb by doubling.
246 pub patience: usize,
247 /// How many partitions one vector may be written into, at most.
248 ///
249 /// One is no replication and is what the index did before this existed. It
250 /// is not the default, and `src/miss.rs` is why.
251 ///
252 /// A vector belongs to the partition whose centroid it is nearest, and on
253 /// some data that is a much weaker statement than it sounds. Measured on a
254 /// million MS-MARCO passage embeddings, only 0.8952 of the true nearest
255 /// neighbours of a query sat in one of the 128 partitions the search reads,
256 /// and the recall the search actually returned was 0.8942, so the whole of
257 /// the miss was neighbours nobody looked at rather than anything the
258 /// estimator did. A vector near the boundary between two partitions is one
259 /// query away from being in the wrong one, and no amount of scanning fixes
260 /// that because the scan never gets there.
261 ///
262 /// So a vector near a boundary goes in both, which is SPANN's answer.
263 /// Raising this raises recall and costs memory and scan time in proportion
264 /// to how many vectors actually qualify, which is what [`Tuning::slack`]
265 /// controls.
266 pub spill: usize,
267 /// How much further than the nearest centroid a vector will still be copied
268 /// into, as a fraction.
269 ///
270 /// A vector goes into every one of its [`Tuning::spill`] nearest partitions
271 /// whose centroid is within `1 + slack` of the nearest one, so zero is no
272 /// replication whatever `spill` says and a large value replicates
273 /// everything into everything. It is a distance ratio rather than a count
274 /// because the thing being asked is whether a vector is genuinely near a
275 /// boundary, and a vector sitting squarely inside its partition should cost
276 /// one copy however large `spill` is.
277 pub slack: f32,
278}
279
280impl Default for Tuning {
281 fn default() -> Tuning {
282 Tuning {
283 posting: 256,
284 probe: 8,
285 rerank: 4,
286 sweep: 4,
287 widen: 8,
288 spill: 4,
289 slack: 0.10,
290 patience: 8,
291 }
292 }
293}
294
295/// The fewest candidates a search will rerank, whatever `k` and `rerank`
296/// multiply out to.
297///
298/// Four times `k` is the right ratio and it is the wrong number when `k` is
299/// small: asking for one answer and reranking four candidates puts the whole
300/// weight of the answer on the estimator getting its top four right, which is
301/// not what the estimator is for. Reranking a few dozen costs a few dozen
302/// squared distances, which is nothing next to the scan that produced them.
303const FLOOR: usize = 32;
304
305/// What decides whether the scan bothers with a member.
306///
307/// A filtered vector search is a recall lottery when the filter runs after the
308/// search: ask for ten English passages, get the best forty by vector, find
309/// three of them are English, and the other seven English passages that were
310/// nearer never had a chance. The fix is to filter inside the scan, so that
311/// only members that can be answers are ranked at all, and that means the thing
312/// the filter reads has to sit next to the codes rather than behind a lookup
313/// into somebody else's table.
314///
315/// So every member carries a `u64` tag, given at insert, and a filter is a
316/// predicate on that tag. What the tag means is the caller's business. A
317/// handful of low cardinality attributes pack into it exactly, one field each,
318/// and the filter is then exact. Anything wider goes through [`Signature`],
319/// which is exact in the direction that matters: it never rejects a member that
320/// should have matched, so the caller's real predicate over the answers still
321/// decides.
322/// A tag that is only a summary needs a second test somewhere, and the place
323/// for it is [`Filter::exact`], which sees the member's id and can go and read
324/// whatever the caller keyed by it. That runs only for members the tag let
325/// through that are also near enough to be ranked, which is why it is allowed to
326/// be the expensive one: an expression over a JSON string is fine there and
327/// would not be fine in the scan.
328pub trait Filter {
329 /// Whether a member with this tag is worth ranking.
330 fn allows(&self, tag: u64) -> bool;
331
332 /// The second test, on the member's id rather than on its tag.
333 ///
334 /// Everything lets everything through by default, because for a filter whose
335 /// tag says the whole truth there is nothing left to ask. Override it when
336 /// the tag is a summary and the real predicate lives in a table of the
337 /// caller's, and keep the tag test as the cheap superset of it: a member the
338 /// tag rejects never reaches here.
339 fn exact(&self, _id: u64) -> bool {
340 true
341 }
342
343 /// Whether this filter can turn members away.
344 ///
345 /// A search that can be turned away has to be ready to look further than
346 /// `probe` partitions, because the nearest ones may not hold `k` members
347 /// that pass, and getting ready to costs something before the first
348 /// partition is read: the probe order has to be built `widen` times longer
349 /// than the search will use if it never widens. Saying no here is how a
350 /// search that cannot be turned away avoids paying for the case that cannot
351 /// happen to it. The default is yes, because a filter that answers wrongly
352 /// here returns short answers rather than slow ones.
353 fn narrowing(&self) -> bool {
354 true
355 }
356}
357
358/// The filter that lets everything through, which is what an unfiltered search
359/// runs.
360#[derive(Debug, Clone, Copy, Default)]
361pub struct Any;
362
363impl Filter for Any {
364 fn allows(&self, _tag: u64) -> bool {
365 true
366 }
367
368 fn narrowing(&self) -> bool {
369 false
370 }
371}
372
373impl<F: Fn(u64) -> bool> Filter for F {
374 fn allows(&self, tag: u64) -> bool {
375 self(tag)
376 }
377}
378
379/// A tag built by setting one bit per attribute value, so that a conjunction of
380/// required values is a subset test.
381///
382/// Superimposed coding, which is old and still the right answer when the test
383/// has to be one instruction on a value that is already in a register. Each
384/// attribute and value pair hashes to one of 64 bits. A member's tag is the
385/// bits for the values it has. A query's tag is the bits for the values it
386/// requires. The member is worth ranking when it has all of the query's bits.
387///
388/// Two different values can land on the same bit, so a member can pass a filter
389/// it does not really match. It can never fail one it does match, which is the
390/// direction that matters: the answers are a superset of the truth and the
391/// caller's own predicate cuts them down, where the other way round would lose
392/// answers silently.
393///
394/// ```
395/// use yo_vector::Signature;
396///
397/// // What a document is tagged with, and what a query asks for.
398/// let doc = Signature::of(&[("lang", "en".as_bytes()), ("topic", "finance".as_bytes())]);
399/// let english = Signature::of(&[("lang", "en".as_bytes())]);
400///
401/// assert!(doc.covers(english));
402/// // The other way round only holds if the two bits happened to collide.
403/// assert!(!english.covers(doc) || english.bits() == doc.bits());
404/// ```
405#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
406pub struct Signature(u64);
407
408impl Signature {
409 /// The signature of a set of attribute and value pairs.
410 #[must_use]
411 pub fn of(values: &[(&str, &[u8])]) -> Signature {
412 let mut got = Signature(0);
413 for (attribute, value) in values {
414 got.insert(attribute, value);
415 }
416 got
417 }
418
419 /// Add one attribute and value pair to what this signature covers.
420 ///
421 /// For a caller that meets the pairs one at a time rather than holding them
422 /// all in a slice, which is what building a tag out of a document's indexed
423 /// fields looks like.
424 pub fn insert(&mut self, attribute: &str, value: &[u8]) {
425 self.insert_bytes(attribute.as_bytes(), value);
426 }
427
428 /// The same for an attribute that is already bytes, which is what a document
429 /// path is.
430 pub fn insert_bytes(&mut self, attribute: &[u8], value: &[u8]) {
431 self.0 |= 1u64 << (hash(attribute, value) % 64);
432 }
433
434 /// The signature as the tag to hand to [`Partitions::insert_tagged`].
435 #[must_use]
436 pub fn bits(self) -> u64 {
437 self.0
438 }
439
440 /// The signature of a tag that came back out of the index.
441 #[must_use]
442 pub fn from_bits(bits: u64) -> Signature {
443 Signature(bits)
444 }
445
446 /// Whether this has every bit `want` has, which is the test the scan runs.
447 #[must_use]
448 pub fn covers(self, want: Signature) -> bool {
449 self.0 & want.0 == want.0
450 }
451}
452
453impl Filter for Signature {
454 fn allows(&self, tag: u64) -> bool {
455 Signature(tag).covers(*self)
456 }
457}
458
459/// FNV over the attribute and then the value, which is small, has no state and
460/// spreads a short value over the whole word well enough to pick a bit.
461fn hash(attribute: &[u8], value: &[u8]) -> u64 {
462 let mut h = 0xcbf2_9ce4_8422_2325u64;
463 for byte in attribute.iter().chain(b":").chain(value) {
464 h ^= u64::from(*byte);
465 h = h.wrapping_mul(0x1000_0000_01b3);
466 }
467 h
468}
469
470/// An answer: a document id and how far it really is, not how far it was
471/// estimated to be.
472#[derive(Debug, Clone, Copy, PartialEq)]
473pub struct Hit {
474 /// The id that was inserted.
475 pub id: u64,
476 /// The exact squared distance, measured against the full precision vector.
477 pub distance: f32,
478}
479
480/// What one search actually read.
481///
482/// [`Tuning::probe`] is a budget rather than a bill, and with
483/// [`Tuning::patience`] set the two are different for most queries. This is the
484/// bill.
485#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
486pub struct Work {
487 /// Partitions read.
488 pub probed: usize,
489 /// Coded members the estimator was run over, which is the number that
490 /// actually tracks the time a search took. Partitions are not the same size
491 /// as each other and a count of them hides that.
492 pub scanned: usize,
493}
494
495/// Where one copy of a member sits, and where the next copy of it is.
496///
497/// A vector is in one posting most of the time and in several when it sits near
498/// the boundary between them, which is what [`Tuning::spill`] is for. So an id
499/// does not map to a place, it maps to a chain of them, threaded through an
500/// arena so that a chain of one costs what the single slot used to cost and a
501/// replicated id costs twelve more bytes per copy rather than an allocation.
502#[derive(Debug, Clone, Copy)]
503struct Place {
504 partition: u32,
505 slot: u32,
506 /// The next copy of the same vector, or [`END`]. Also the free list.
507 next: u32,
508}
509
510/// The end of a chain, and the empty free list.
511const END: u32 = u32::MAX;
512
513/// One partition's members: the ids, their codes end to end, and what each code
514/// needs beside it.
515#[derive(Debug, Default)]
516struct Posting {
517 ids: Vec<u64>,
518 /// One tag per member, in the same order, which is what a filter meets.
519 ///
520 /// Beside the ids rather than behind a pointer, because the whole point is
521 /// that the scan can skip a member without touching anything that is not
522 /// already in cache.
523 tags: Vec<u64>,
524 codes: Vec<u8>,
525 meta: Vec<Coded>,
526 /// The size at which a split was tried and found there was no cut, which
527 /// happens when every member is the same vector. It is not tried again
528 /// until the posting has grown past it.
529 stuck: usize,
530}
531
532impl Posting {
533 fn len(&self) -> usize {
534 self.ids.len()
535 }
536}
537
538/// A collection of vectors, quantised, partitioned, and updated in place.
539#[derive(Debug)]
540pub struct Partitions {
541 quant: Quantizer,
542 tuning: Tuning,
543 /// The centroids, already rotated, `dim` floats each end to end.
544 centroids: Vec<f32>,
545 postings: Vec<Posting>,
546 /// The head of every id's chain of placements, which is what makes a delete
547 /// a constant time operation rather than a search.
548 at: HashMap<u64, u32>,
549 /// The placements themselves, and the free list through their `next`.
550 ///
551 /// One arena rather than a list per id, because most ids have exactly one
552 /// placement and a `Vec` each would be a million allocations on a million
553 /// vectors to hold one entry apiece.
554 places: Vec<Place>,
555 free: u32,
556 /// The index over the centroids. See [`crate::coarse`].
557 coarse: Coarse,
558 /// The shortlist a placement fills in, kept here so that placing a vector
559 /// does not allocate.
560 scratch: Vec<u32>,
561 /// One member's copies, and the partitions an insert is about to spill
562 /// into, kept for the same reason `scratch` is.
563 spare: Vec<Place>,
564 spill: Vec<(usize, f32)>,
565 /// Partitions that may be over the split threshold, and partitions that may
566 /// be under the merge threshold.
567 ///
568 /// Deciding what to maintain next used to be two passes over every
569 /// partition, once for the largest and once for the smallest, and it ran
570 /// once per vector inserted. That is linear in the partition count and the
571 /// partition count grows with the collection, which is the same quadratic
572 /// the coarse layer was built to remove, just hiding somewhere else. By 1.6
573 /// million vectors it was a fifth of an ingest and it was doing nothing at
574 /// all almost every time it ran.
575 ///
576 /// A partition can only cross a threshold when its own length changes, and
577 /// there are five places a length changes, so the crossings can be recorded
578 /// as they happen instead of looked for afterwards. These two lists hold
579 /// every partition that qualifies and usually nothing else. They are
580 /// allowed to hold stale entries, because [`Partitions::job`] checks the
581 /// real length before it returns anything and drops what no longer
582 /// qualifies, and they are kept in partition order so that a tie is broken
583 /// the same way the two passes broke it.
584 big: Vec<u32>,
585 small: Vec<u32>,
586}
587
588impl Partitions {
589 /// An empty collection of `dim` dimensional vectors.
590 ///
591 /// The first vector inserted becomes the first centroid, and the index
592 /// grows by splitting from there, so there is no build step and no moment
593 /// where the shape of the collection has to be known in advance.
594 ///
595 /// # Panics
596 ///
597 /// If `dim` is zero.
598 #[must_use]
599 pub fn new(dim: usize, bits: Bits, seed: u64, tuning: Tuning) -> Partitions {
600 Partitions {
601 quant: Quantizer::new(dim, bits, seed),
602 tuning,
603 centroids: Vec::new(),
604 postings: Vec::new(),
605 at: HashMap::new(),
606 places: Vec::new(),
607 free: END,
608 coarse: Coarse::default(),
609 scratch: Vec::new(),
610 spare: Vec::new(),
611 spill: Vec::new(),
612 big: Vec::new(),
613 small: Vec::new(),
614 }
615 }
616
617 /// How many coordinates a vector here has.
618 #[must_use]
619 pub fn dim(&self) -> usize {
620 self.quant.dim()
621 }
622
623 /// How many vectors are in the collection.
624 #[must_use]
625 pub fn len(&self) -> usize {
626 self.at.len()
627 }
628
629 /// Whether there are none.
630 #[must_use]
631 pub fn is_empty(&self) -> bool {
632 self.at.is_empty()
633 }
634
635 /// How many partitions the collection has grown to.
636 #[must_use]
637 pub fn partitions(&self) -> usize {
638 self.postings.len()
639 }
640
641 /// How many coded members the postings hold between them.
642 ///
643 /// The same as [`Partitions::len`] until [`Tuning::spill`] puts a vector
644 /// near a boundary into more than one partition, and the ratio of the two
645 /// is what replication is costing in memory and in scan time.
646 #[must_use]
647 pub fn entries(&self) -> usize {
648 self.postings.iter().map(Posting::len).sum()
649 }
650
651 /// The knobs.
652 #[must_use]
653 pub fn tuning(&self) -> Tuning {
654 self.tuning
655 }
656
657 /// Change the knobs on a collection that already has vectors in it.
658 ///
659 /// [`Tuning::probe`], [`Tuning::rerank`] and [`Tuning::widen`] are read by
660 /// each search, so they take effect on the next one. That is what makes a
661 /// recall against latency curve measurable on one built index rather than on
662 /// one built per row, and it is what `EF_RUNTIME` means to a client that
663 /// thinks it is talking to a graph.
664 ///
665 /// [`Tuning::posting`] and [`Tuning::sweep`] are what maintenance aims at,
666 /// so lowering `posting` does not split anything by itself. The partitions
667 /// move towards the new size as [`Partitions::maintain`] gets called, which
668 /// is the same way they got to the old one.
669 pub fn retune(&mut self, tuning: Tuning) {
670 self.tuning = tuning;
671 }
672
673 /// The quantiser, whose seed and width a catalogue entry has to record.
674 #[must_use]
675 pub fn quantizer(&self) -> &Quantizer {
676 &self.quant
677 }
678
679 /// How many bytes the codes take, which is the searchable size of the
680 /// collection and the number the 32x claim is about.
681 #[must_use]
682 pub fn code_bytes(&self) -> usize {
683 self.postings.iter().map(|p| p.codes.len()).sum()
684 }
685
686 /// Put a vector in, replacing whatever was under `id`.
687 ///
688 /// Two appends and no locks: the code goes on the end of the nearest
689 /// partition's posting, and the caller puts the full precision vector in the
690 /// log. Nothing else in the index is touched, which is the difference
691 /// between this and a graph.
692 ///
693 /// # Panics
694 ///
695 /// If `v` is not [`Partitions::dim`] long.
696 pub fn insert(&mut self, id: u64, v: &[f32]) {
697 self.insert_tagged(id, v, 0);
698 }
699
700 /// The same, with the tag a filter will meet in the scan.
701 ///
702 /// See [`Filter`] for what a tag is and [`Signature`] for the encoding to
703 /// reach for when the attributes do not fit in one exactly.
704 ///
705 /// # Panics
706 ///
707 /// If `v` is not [`Partitions::dim`] long.
708 pub fn insert_tagged(&mut self, id: u64, v: &[f32], tag: u64) {
709 assert_eq!(
710 v.len(),
711 self.dim(),
712 "this collection holds {} dimensional vectors and was handed {}",
713 self.dim(),
714 v.len()
715 );
716 self.remove(id);
717 let x = self.quant.rotate(v);
718 if self.postings.is_empty() {
719 // The first vector is the first centroid. There is nothing to
720 // average it with yet, and the first split is what starts the
721 // centroids being means rather than members.
722 let p = self.add_partition(&x);
723 self.place(p, id, tag, &x);
724 return;
725 }
726 let mut into = core::mem::take(&mut self.spill);
727 self.spill_into(&x, &mut into);
728 for &(p, _) in &into {
729 self.place(p, id, tag, &x);
730 }
731 self.spill = into;
732 }
733
734 /// The partitions a vector goes into, nearest first.
735 ///
736 /// The first is the one it belongs to and there is always exactly one of
737 /// those. The rest are the boundary copies [`Tuning::spill`] is about: every
738 /// further partition whose centroid is within [`Tuning::slack`] of the
739 /// nearest one, up to `spill` of them in total.
740 ///
741 /// # The rule that is not here
742 ///
743 /// SPANN has a third condition, and it was written, measured and taken out
744 /// again. It drops a candidate if some partition already chosen is nearer to
745 /// it than the vector is, on the grounds that a candidate on the far side of
746 /// one already taken adds a copy in a direction that is already covered.
747 /// That is the rule that keeps SPANN's replication factor down.
748 ///
749 /// It rejects every candidate there is at a thousand dimensions. On two
750 /// hundred thousand generated 1024 dimensional vectors in 528 partitions,
751 /// the copy rate with the rule in is 1.0000 at every setting of `spill` and
752 /// `slack` that was tried, and without it 2.85 at a `spill` of 4 and 5.08 at
753 /// 8. A million MS-MARCO passages say the same thing from the other end:
754 /// 1.000 copies a vector at `spill` 4 and 1.001 at `spill` 8 with `slack` at
755 /// 0.60, which is a feature that is switched on and doing nothing.
756 ///
757 /// The reason is the one [`coarse`](crate::coarse) already ran into.
758 /// Distances concentrate, and a centroid is the mean of a few hundred
759 /// members so it sits well inside a cloud whose radius is most of the
760 /// distance to the next centroid. Two neighbouring centroids are therefore
761 /// much closer to each other than any of their members is to either, the
762 /// condition holds for every pair, and nothing is ever copied. Keeping a
763 /// rule that fires on nothing would have left the whole feature switched on
764 /// and inert, which is worse than not having it.
765 ///
766 /// What is left holds the replication factor down instead: `spill` caps it
767 /// outright and `slack` cuts off candidates that are not really boundary
768 /// cases. On this data `spill` is what binds, because `slack` at 0.15 and at
769 /// 0.35 give copy rates of 2.8500 and 2.8452, which is the same
770 /// concentration seen from the other side.
771 fn spill_into(&mut self, x: &[f32], into: &mut Vec<(usize, f32)>) {
772 into.clear();
773 let dim = self.dim();
774 let want = self.tuning.spill.max(1);
775 // The coarse layer's shortlist rather than every centroid, which is what
776 // keeps an insert from costing what a search costs. It is at least 256
777 // partitions wide, so the nearest handful of them are in there.
778 let mut short = core::mem::take(&mut self.scratch);
779 if want == 1 || self.tuning.slack <= 0.0 {
780 let p = self.roughly_nearest(x, &mut short);
781 self.scratch = short;
782 into.push((p, 0.0));
783 return;
784 }
785 self.coarse.shortlist(x, dim, &mut short);
786 let mut near: Vec<(usize, f32)> = if short.is_empty() {
787 (0..self.postings.len())
788 .map(|p| (p, sqdist(x, self.centroid(p))))
789 .collect()
790 } else {
791 short
792 .iter()
793 .map(|&p| (p as usize, sqdist(x, self.centroid(p as usize))))
794 .collect()
795 };
796 self.scratch = short;
797 near.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
798 let Some(&(first, best)) = near.first() else {
799 return;
800 };
801 into.push((first, best));
802 // Squared distances throughout, so the ratio on distances is the square
803 // of it here. Comparing the squares directly and skipping two square
804 // roots per candidate is worth it on a path that runs once per insert.
805 let ceiling = best * (1.0 + self.tuning.slack) * (1.0 + self.tuning.slack);
806 for &(q, d) in near.iter().skip(1) {
807 if into.len() >= want {
808 break;
809 }
810 if d > ceiling {
811 break;
812 }
813 into.push((q, d));
814 }
815 }
816
817 /// The tag `id` was inserted with, if it is still here.
818 #[must_use]
819 pub fn tag(&self, id: u64) -> Option<u64> {
820 // Any copy will do. Every copy of a member carries the same tag, which
821 // is what [`Partitions::retag`] is for.
822 let at = self.any_place(id)?;
823 Some(self.postings[at.partition as usize].tags[at.slot as usize])
824 }
825
826 /// Change the tag `id` carries, saying whether it was there.
827 ///
828 /// The tag sits beside the code and nothing about the placement depends on
829 /// it, so this is one write and no maintenance. That is what makes it
830 /// affordable to recompute every tag in a collection when the thing the tag
831 /// summarises changes, which for a document index is a field being indexed
832 /// or stopping being indexed.
833 pub fn retag(&mut self, id: u64, tag: u64) -> bool {
834 let mut walk = self.at.get(&id).copied().unwrap_or(END);
835 let mut found = false;
836 while walk != END {
837 let place = self.places[walk as usize];
838 self.postings[place.partition as usize].tags[place.slot as usize] = tag;
839 found = true;
840 walk = place.next;
841 }
842 found
843 }
844
845 /// Take a vector out, saying whether it was there.
846 ///
847 /// The last member of the posting moves into the hole. There is no
848 /// tombstone, so there is nothing to accumulate and nothing to compact.
849 pub fn remove(&mut self, id: u64) -> bool {
850 let mut copies = core::mem::take(&mut self.spare);
851 self.every_place(id, &mut copies);
852 if copies.is_empty() {
853 self.spare = copies;
854 return false;
855 }
856 self.detach_all(id);
857 // Highest slot first inside a partition, because pulling a member moves
858 // the last one into its slot, and a copy of this same id sitting at a
859 // higher slot in the same posting would have its recorded slot go stale.
860 // There is at most one copy per partition so this only matters across
861 // them, but the order costs nothing and the alternative is a rule that
862 // has to stay true.
863 copies.sort_unstable_by_key(|c| core::cmp::Reverse(c.slot));
864 for copy in &copies {
865 let p = copy.partition as usize;
866 let s = copy.slot as usize;
867 if let Some(moved) = self.pull(p, s) {
868 self.reslot(moved, p, s);
869 }
870 }
871 self.spare = copies;
872 true
873 }
874
875 /// Whether `id` is in the collection.
876 #[must_use]
877 pub fn contains(&self, id: u64) -> bool {
878 self.at.contains_key(&id)
879 }
880
881 // -- what an image is made of -------------------------------------------
882 //
883 // [`crate::image`] writes an index down and reads it back, and it lives in
884 // its own file because the layout it writes is the format's business rather
885 // than the index's. These four are the seam between the two: everything
886 // above is private on purpose and none of it is worth making public just so
887 // that a sibling module can copy it into a buffer.
888
889 /// Every centroid, already rotated, `dim` floats each end to end.
890 pub(crate) fn all_centroids(&self) -> &[f32] {
891 &self.centroids
892 }
893
894 /// One partition's four parallel arrays, and the size at which its last
895 /// split gave up.
896 pub(crate) fn posting_parts(&self, p: usize) -> (&[u64], &[u64], &[u8], &[Coded], usize) {
897 let posting = &self.postings[p];
898 (
899 &posting.ids,
900 &posting.tags,
901 &posting.codes,
902 &posting.meta,
903 posting.stuck,
904 )
905 }
906
907 /// Put a whole partition back, centroid and members together.
908 ///
909 /// The centroid goes on the end of the run and the members go into a new
910 /// posting, so partitions come back in the order they were written and an
911 /// id keeps the partition number it had. Nothing is requantised and nothing
912 /// is measured: an image holds the codes, and recomputing them from the
913 /// vectors would be the rebuild this whole index exists to not do.
914 ///
915 /// # Errors
916 ///
917 /// [`Code::Corrupt`] if the four arrays do not describe the same members or
918 /// if an id is already in the index.
919 pub(crate) fn absorb(
920 &mut self,
921 centroid: &[f32],
922 ids: Vec<u64>,
923 tags: Vec<u64>,
924 codes: Vec<u8>,
925 meta: Vec<Coded>,
926 stuck: usize,
927 ) -> Result<()> {
928 let width = self.quant.code_bytes();
929 if centroid.len() != self.dim()
930 || tags.len() != ids.len()
931 || meta.len() != ids.len()
932 || codes.len() != ids.len() * width
933 {
934 return Err(Error::new(
935 Code::Corrupt,
936 "the parts of a partition do not describe the same members",
937 )
938 .with_detail(format!(
939 "centroid={} ids={} tags={} codes={} meta={}",
940 centroid.len(),
941 ids.len(),
942 tags.len(),
943 codes.len(),
944 meta.len()
945 )));
946 }
947 let p = self.postings.len();
948 for (slot, &id) in ids.iter().enumerate() {
949 // An id in two partitions is a replicated member and is what an
950 // image of a spilled collection looks like. An id twice in one
951 // partition is not, and `attach` is where that is caught, because
952 // it is the shape a delete cannot undo.
953 if !self.attach(id, p, slot) {
954 return Err(Error::new(
955 Code::Corrupt,
956 "an id is twice in one partition of an image",
957 )
958 .with_detail(format!("id={id} partition={p}")));
959 }
960 }
961 self.centroids.extend_from_slice(centroid);
962 self.postings.push(Posting {
963 ids,
964 tags,
965 codes,
966 meta,
967 stuck,
968 });
969 Ok(())
970 }
971
972 /// Say that a load is over, so the derived parts can be built once.
973 ///
974 /// The coarse layer and the two maintenance candidate lists are the whole of
975 /// what an image does not carry, because both are decided by the centroids
976 /// and the posting lengths that it does carry. Building them here is one
977 /// pass rather than the running updates the insert path makes, which is the
978 /// difference between a load being linear and being quadratic.
979 pub(crate) fn finish_image(&mut self) {
980 let dim = self.quant.dim();
981 let n = self.postings.len();
982 self.coarse.rebuild(&self.centroids, dim, n);
983 self.big.clear();
984 self.small.clear();
985 for p in 0..n {
986 if self.over(p) {
987 self.big.push(p as u32);
988 }
989 if self.under(p) {
990 self.small.push(p as u32);
991 }
992 }
993 }
994
995 /// The `k` nearest vectors to `q`, measured exactly.
996 ///
997 /// The codes pick the candidates and the log settles the order, so the
998 /// answer is as exact as brute force whenever the candidates contained the
999 /// truth, and the recall table is about how often they do.
1000 ///
1001 /// # Panics
1002 ///
1003 /// If `q` is not [`Partitions::dim`] long.
1004 #[must_use]
1005 pub fn search(&self, q: &[f32], k: usize, vectors: &impl Vectors) -> Vec<Hit> {
1006 self.search_where(q, k, &Any, vectors)
1007 }
1008
1009 /// The `k` nearest vectors to `q` that a filter allows.
1010 ///
1011 /// The filter runs inside the scan, on the tag that sits next to the code,
1012 /// so a member the filter rejects is never ranked and never takes a place
1013 /// that an answer should have had. Filtering afterwards instead is what
1014 /// makes a filtered vector search a lottery, and the more selective the
1015 /// filter the worse a lottery it is.
1016 ///
1017 /// A selective filter also means the nearest few partitions may not hold `k`
1018 /// members that pass, so the scan keeps going into further partitions until
1019 /// it has enough or until it has spent [`Tuning::widen`]. A filter that
1020 /// matches almost nothing returns fewer answers rather than reading the
1021 /// whole collection, which is the trade every engine makes here and is worth
1022 /// saying out loud.
1023 ///
1024 /// # Panics
1025 ///
1026 /// If `q` is not [`Partitions::dim`] long.
1027 #[must_use]
1028 pub fn search_where(
1029 &self,
1030 q: &[f32],
1031 k: usize,
1032 filter: &impl Filter,
1033 vectors: &impl Vectors,
1034 ) -> Vec<Hit> {
1035 self.search_costed(q, k, filter, vectors).0
1036 }
1037
1038 /// The same again, and what the scan behind it cost.
1039 ///
1040 /// See [`Work`]. Worth having in front of a caller rather than behind a
1041 /// feature flag, because with [`Tuning::patience`] set the cost of a search
1042 /// is a property of the query and not of the settings, and a tuner that
1043 /// cannot see it is guessing.
1044 ///
1045 /// # Panics
1046 ///
1047 /// If `q` is not [`Partitions::dim`] long.
1048 #[must_use]
1049 pub fn search_costed(
1050 &self,
1051 q: &[f32],
1052 k: usize,
1053 filter: &impl Filter,
1054 vectors: &impl Vectors,
1055 ) -> (Vec<Hit>, Work) {
1056 if k == 0 {
1057 return (Vec::new(), Work::default());
1058 }
1059 let (candidates, work) =
1060 self.candidates_costed(q, (k * self.tuning.rerank).max(FLOOR), filter);
1061 let mut buf = vec![0.0f32; self.dim()];
1062 let mut hits = Vec::with_capacity(candidates.len());
1063 for (id, _) in candidates {
1064 if vectors.get(id, &mut buf) {
1065 hits.push(Hit {
1066 id,
1067 distance: sqdist(q, &buf),
1068 });
1069 }
1070 }
1071 hits.sort_by(|a, b| a.distance.total_cmp(&b.distance));
1072 hits.truncate(k);
1073 (hits, work)
1074 }
1075
1076 /// The `want` best candidates by the estimator, without rerank.
1077 ///
1078 /// This is what a filter will eventually push into, and it is what the
1079 /// recall of the codes alone is measured on.
1080 ///
1081 /// # Panics
1082 ///
1083 /// If `q` is not [`Partitions::dim`] long.
1084 #[must_use]
1085 pub fn candidates(&self, q: &[f32], want: usize) -> Vec<(u64, f32)> {
1086 self.candidates_where(q, want, &Any)
1087 }
1088
1089 /// The same, with the filter run inside the scan.
1090 ///
1091 /// # Panics
1092 ///
1093 /// If `q` is not [`Partitions::dim`] long.
1094 #[must_use]
1095 pub fn candidates_where(
1096 &self,
1097 q: &[f32],
1098 want: usize,
1099 filter: &impl Filter,
1100 ) -> Vec<(u64, f32)> {
1101 self.candidates_costed(q, want, filter).0
1102 }
1103
1104 /// The same again, and what reading them cost.
1105 ///
1106 /// The cost is here because [`Tuning::patience`] makes it vary from one
1107 /// query to the next, and a knob whose whole point is that different queries
1108 /// pay different amounts is not one anybody can set without being able to see
1109 /// what it did. It is also the honest way to compare two settings: recall
1110 /// against partitions actually read, rather than recall against the budget
1111 /// neither of them spent.
1112 ///
1113 /// # Panics
1114 ///
1115 /// If `q` is not [`Partitions::dim`] long.
1116 #[must_use]
1117 pub fn candidates_costed(
1118 &self,
1119 q: &[f32],
1120 want: usize,
1121 filter: &impl Filter,
1122 ) -> (Vec<(u64, f32)>, Work) {
1123 assert_eq!(
1124 q.len(),
1125 self.dim(),
1126 "this collection holds {} dimensional vectors and was handed {}",
1127 self.dim(),
1128 q.len()
1129 );
1130 if want == 0 || self.postings.is_empty() {
1131 return (Vec::new(), Work::default());
1132 }
1133 // Rotated once here and never again, which is what lets a search probe
1134 // tens of partitions without paying for tens of rotations.
1135 let u = self.quant.rotate(q);
1136 let mut best = Bounded::new(want);
1137 // One buffer for the whole search rather than one per partition, and
1138 // grown rather than cleared, because every partition after the first
1139 // wants the same room the one before it did.
1140 let mut scores: Vec<f32> = Vec::new();
1141 // A filter that turns nothing away can never widen, so it does not have
1142 // to pay for a probe order it will not read.
1143 let reach = if filter.narrowing() {
1144 self.tuning.probe.saturating_mul(self.tuning.widen.max(1))
1145 } else {
1146 self.tuning.probe
1147 };
1148 let mut work = Work::default();
1149 // How many partitions in a row have gone by without one of their members
1150 // being good enough to displace an answer. See [`Tuning::patience`].
1151 let mut quiet = 0;
1152 for (n, p) in self.near_partitions(&u, reach).into_iter().enumerate() {
1153 // Two reasons to stop, and both of them need enough answers in hand
1154 // first. Past the partitions an unfiltered search would have read,
1155 // keep going only while there is still not enough to answer with; an
1156 // unfiltered search never gets here, because the first `probe`
1157 // partitions of a collection worth probing hold more than `want`.
1158 // Inside them, stop once the last few have added nothing.
1159 if best.full()
1160 && (n >= self.tuning.probe
1161 || (self.tuning.patience > 0 && quiet >= self.tuning.patience))
1162 {
1163 break;
1164 }
1165 let prepared = self.quant.query_rotated(&u, self.centroid(p));
1166 let posting = &self.postings[p];
1167 let held = posting.ids.len();
1168 if scores.len() < held {
1169 scores.resize(held, 0.0);
1170 }
1171 // The whole posting at once, so the estimator's inner loops know
1172 // how wide a code is. Then a second pass, which for most members is
1173 // one comparison against the worst answer so far and no more, and
1174 // which does not read the id or the tag of a member that lost.
1175 prepared.scan(&posting.codes, &posting.meta, &mut scores[..held]);
1176 work.probed += 1;
1177 work.scanned += held;
1178 let mut took = 0;
1179 for (i, &at) in scores[..held].iter().enumerate() {
1180 if !best.wants(at) {
1181 continue;
1182 }
1183 if !filter.allows(posting.tags[i]) {
1184 continue;
1185 }
1186 if !filter.exact(posting.ids[i]) {
1187 continue;
1188 }
1189 best.put(posting.ids[i], at);
1190 took += 1;
1191 }
1192 // A partition that displaced something buys the search its patience
1193 // back, because a run of empty ones followed by a good one is the
1194 // shape of a query whose neighbourhood is spread out rather than a
1195 // query that has finished.
1196 quiet = if took == 0 { quiet + 1 } else { 0 };
1197 }
1198 let mut out = best.sorted();
1199 // A replicated member is in more than one posting and a search can read
1200 // more than one of them, so the same id can be ranked twice, with two
1201 // different estimates because each copy is coded against its own
1202 // centroid. The near duplicates are not adjacent for that reason, so
1203 // this is a pass with a set rather than a `dedup`.
1204 //
1205 // Only when there is replication to undo. With `spill` at one there can
1206 // be no duplicate, and a search that pays for proving it every time is
1207 // charging every collection for a feature some of them do not use.
1208 if self.tuning.spill > 1 {
1209 let mut seen = HashSet::with_capacity(out.len());
1210 out.retain(|&(id, _)| seen.insert(id));
1211 }
1212 (out, work)
1213 }
1214
1215 /// Whether there is a split or a merge waiting.
1216 #[must_use]
1217 pub fn needs_maintenance(&self) -> bool {
1218 // The same two questions [`Partitions::job`] asks, without the pruning,
1219 // so that asking does not need the collection mutably.
1220 self.big.iter().any(|&p| {
1221 let p = p as usize;
1222 self.over(p) && self.postings[p].len() > self.postings[p].stuck
1223 }) || (self.postings.len() > 1 && self.small.iter().any(|&p| self.under(p as usize)))
1224 }
1225
1226 /// Whether partition `p` is big enough to split. False for an index that is
1227 /// no longer there, which is what a stale candidate looks like.
1228 fn over(&self, p: usize) -> bool {
1229 self.postings
1230 .get(p)
1231 .is_some_and(|posting| posting.len() > self.tuning.posting * 2)
1232 }
1233
1234 /// Whether partition `p` is small enough to merge away.
1235 fn under(&self, p: usize) -> bool {
1236 self.postings
1237 .get(p)
1238 .is_some_and(|posting| posting.len() * 4 < self.tuning.posting)
1239 }
1240
1241 /// Do bounded maintenance, and say how many vectors it looked at.
1242 ///
1243 /// `budget` is in vectors touched rather than in time, because time is not
1244 /// something a storage engine gets to measure cheaply and a vector is the
1245 /// unit all of this work is actually made of. Call it from a maintenance
1246 /// slice until it returns less than the budget, which means there was
1247 /// nothing left to do.
1248 pub fn maintain(&mut self, vectors: &impl Vectors, budget: usize) -> usize {
1249 let mut done = 0;
1250 while done < budget {
1251 let Some(job) = self.job() else { break };
1252 done += match job {
1253 Job::Split(p) => self.split(p, vectors),
1254 Job::Merge(p) => self.merge(p, vectors),
1255 };
1256 }
1257 done
1258 }
1259
1260 /// Note that partition `p`'s length has changed, so it may have crossed a
1261 /// threshold in either direction.
1262 ///
1263 /// Both lists are checked, rather than only the one the direction of the
1264 /// change could have reached, because every caller would otherwise have to
1265 /// know which way it moved the length and one of them moves it both ways.
1266 /// They are short enough that looking is free.
1267 fn note(&mut self, p: usize) {
1268 let (over, under) = (self.over(p), self.under(p));
1269 let p = p as u32;
1270 if over && !self.big.contains(&p) {
1271 self.big.push(p);
1272 }
1273 if under && !self.small.contains(&p) {
1274 self.small.push(p);
1275 }
1276 }
1277
1278 /// The next thing worth doing, biggest problem first.
1279 ///
1280 /// This used to walk every partition twice. It now walks the two candidate
1281 /// lists, which hold every partition that qualifies and are usually empty,
1282 /// and drops the entries that have stopped qualifying on the way past. The
1283 /// answer is the same one the walk gave, including which partition is
1284 /// picked when two are the same size, because the lists are in partition
1285 /// order and a maximum takes the last of equals where a minimum takes the
1286 /// first.
1287 fn job(&mut self) -> Option<Job> {
1288 let mut big = std::mem::take(&mut self.big);
1289 big.retain(|&p| self.over(p as usize));
1290 big.sort_unstable();
1291 let split = big
1292 .iter()
1293 .map(|&p| p as usize)
1294 .filter(|&p| self.postings[p].len() > self.postings[p].stuck)
1295 .max_by_key(|&p| self.postings[p].len());
1296 self.big = big;
1297 if let Some(split) = split {
1298 return Some(Job::Split(split));
1299 }
1300
1301 if self.postings.len() > 1 {
1302 let mut small = std::mem::take(&mut self.small);
1303 small.retain(|&p| self.under(p as usize));
1304 small.sort_unstable();
1305 let merge = small
1306 .iter()
1307 .map(|&p| p as usize)
1308 .min_by_key(|&p| self.postings[p].len());
1309 self.small = small;
1310 if let Some(merge) = merge {
1311 return Some(Job::Merge(merge));
1312 }
1313 }
1314 None
1315 }
1316
1317 /// Cut a partition in two by two means over its own members, then sweep the
1318 /// neighbours for anything that should have come along.
1319 fn split(&mut self, p: usize, vectors: &impl Vectors) -> usize {
1320 let (members, xs) = self.take(p, vectors);
1321 let dim = self.dim();
1322 if members.len() < 2 {
1323 for (i, m) in members.iter().enumerate() {
1324 self.place(p, m.id, m.tag, &xs[i * dim..(i + 1) * dim]);
1325 }
1326 return members.len();
1327 }
1328 let (a, b) = two_means(&xs, dim);
1329 let sides: Vec<bool> = (0..members.len())
1330 .map(|i| {
1331 sqdist(&xs[i * dim..(i + 1) * dim], &a) <= sqdist(&xs[i * dim..(i + 1) * dim], &b)
1332 })
1333 .collect();
1334 // A cut that leaves one side small enough to merge straight back has not
1335 // divided anything, and taking it is worse than leaving the posting
1336 // alone. A thousand copies of the same vector is the obvious case: two
1337 // means sees one point, everything lands on one side, and there is no
1338 // cut at all. The case that actually bit is milder and does not look
1339 // like a problem from inside the split. One outlier in an otherwise
1340 // round cloud gets cut off on its own, `merge` sees a partition of one,
1341 // hands the outlier back to its nearest centroid, which is the partition
1342 // it just came out of, and that partition is over the limit again. Split
1343 // and merge then take turns for as long as anyone is willing to call
1344 // `maintain`, at a full budget a call, on a collection that is not
1345 // changing. So both cases put the members back and do not come back
1346 // until the posting has doubled, which costs a collection that really is
1347 // all one vector a re-encode a logarithmic number of times rather than
1348 // once per insert.
1349 let small = sides.iter().filter(|&&s| s).count();
1350 let small = small.min(members.len() - small);
1351 if small == 0 || small * 4 < self.tuning.posting {
1352 for (i, m) in members.iter().enumerate() {
1353 self.place(p, m.id, m.tag, &xs[i * dim..(i + 1) * dim]);
1354 }
1355 self.postings[p].stuck = members.len() * 2;
1356 return members.len();
1357 }
1358 self.centroids[p * dim..(p + 1) * dim].copy_from_slice(&a);
1359 self.coarse.moved(p, &a, dim);
1360 let q = self.add_partition(&b);
1361 for (i, m) in members.iter().enumerate() {
1362 let to = if sides[i] { p } else { q };
1363 self.place(to, m.id, m.tag, &xs[i * dim..(i + 1) * dim]);
1364 }
1365 members.len() + self.sweep(&[p, q], vectors)
1366 }
1367
1368 /// Hand a partition's members to whoever is nearest now, and drop it.
1369 fn merge(&mut self, p: usize, vectors: &impl Vectors) -> usize {
1370 let (members, xs) = self.take(p, vectors);
1371 let dim = self.dim();
1372 self.drop_partition(p);
1373 // The same approximate lookup an insert uses, and for the same reason.
1374 // An exact one is a distance to every centroid in the collection, once
1375 // per member of the partition being emptied, and a merge is the one
1376 // maintenance job whose whole cost is that lookup. Measured at 1024
1377 // dimensions and 4849 partitions it is most of an ingest.
1378 let mut short = core::mem::take(&mut self.scratch);
1379 for (i, m) in members.iter().enumerate() {
1380 let x = &xs[i * dim..(i + 1) * dim];
1381 let to = self.roughly_nearest(x, &mut short);
1382 self.place(to, m.id, m.tag, x);
1383 }
1384 self.scratch = short;
1385 members.len()
1386 }
1387
1388 /// LIRE: after the centroids move, anything nearby that is now filed under
1389 /// the wrong one gets moved.
1390 ///
1391 /// Only the partitions near the ones that just changed are looked at,
1392 /// because those are the only ones whose members can have a new nearest
1393 /// centroid, and looking at all of them would be the rebuild this index
1394 /// exists to avoid.
1395 ///
1396 /// # Why a member is only measured against what changed
1397 ///
1398 /// Every member is already filed under the centroid it was nearest to, and a
1399 /// split moves one centroid and adds one. Nothing else moved, so for a member
1400 /// of some other partition the nearest of all the centroids that did not
1401 /// change is still the one it is already under, and the only way it can have
1402 /// a new answer is if one of the two new centroids beats that. That is a
1403 /// comparison against two, not a search over all of them.
1404 ///
1405 /// This is not a shortcut, it is what LIRE says, and getting it wrong is
1406 /// expensive in a way that is easy to miss. A sweep after a split walks about
1407 /// four partitions' worth of members, and a split happens every posting's
1408 /// worth of inserts, so a full centroid scan per member works out at several
1409 /// scans of every centroid in the collection per vector inserted. That is the
1410 /// whole ingest cost at any size worth talking about: measured on 128
1411 /// dimensional vectors it was 74 thousand a second at twelve thousand vectors
1412 /// and 13 thousand at two hundred thousand, with maintenance three quarters
1413 /// of it, and `examples/ingest.rs` is the harness that says so.
1414 fn sweep(&mut self, changed: &[usize], vectors: &impl Vectors) -> usize {
1415 let dim = self.dim();
1416 let mut look: Vec<usize> = Vec::new();
1417 for &p in changed {
1418 let centre = self.centroid(p).to_vec();
1419 for q in self.roughly_near_partitions(¢re, self.tuning.sweep) {
1420 if !changed.contains(&q) && !look.contains(&q) {
1421 look.push(q);
1422 }
1423 }
1424 }
1425 // Copied out because placing a member borrows the index, and safe to
1426 // copy because nothing below here moves a centroid: `place` appends a
1427 // code to a posting and leaves the centroids alone.
1428 let fresh: Vec<(usize, Vec<f32>)> = changed
1429 .iter()
1430 .map(|&p| (p, self.centroid(p).to_vec()))
1431 .collect();
1432 let mut seen = 0;
1433 let mut buf = vec![0.0f32; dim];
1434 for p in look {
1435 let here = self.centroid(p).to_vec();
1436 // Backwards, because taking a member out moves the last one into
1437 // its slot and a backwards walk never steps over the one that moved.
1438 for i in (0..self.postings[p].len()).rev() {
1439 seen += 1;
1440 let id = self.postings[p].ids[i];
1441 let tag = self.postings[p].tags[i];
1442 if !vectors.get(id, &mut buf) {
1443 self.pull_and_forget(p, i);
1444 continue;
1445 }
1446 let x = self.quant.rotate(&buf);
1447 let mut best = (p, sqdist(&x, &here));
1448 for (q, centre) in &fresh {
1449 let d = sqdist(&x, centre);
1450 if d < best.1 {
1451 best = (*q, d);
1452 }
1453 }
1454 if best.0 != p {
1455 self.pull_and_forget(p, i);
1456 self.place(best.0, id, tag, &x);
1457 }
1458 }
1459 }
1460 seen
1461 }
1462
1463 /// Empty a partition out, handing back its members and their rotated
1464 /// vectors. Ids the source has forgotten are dropped.
1465 fn take(&mut self, p: usize, vectors: &impl Vectors) -> (Vec<Member>, Vec<f32>) {
1466 let dim = self.dim();
1467 let ids = std::mem::take(&mut self.postings[p].ids);
1468 let tags = std::mem::take(&mut self.postings[p].tags);
1469 self.postings[p].codes.clear();
1470 self.postings[p].meta.clear();
1471 self.note(p);
1472 let mut kept = Vec::with_capacity(ids.len());
1473 let mut xs = Vec::with_capacity(ids.len() * dim);
1474 let mut buf = vec![0.0f32; dim];
1475 for (id, tag) in ids.into_iter().zip(tags) {
1476 // Only this partition's copy. A member replicated into a partition
1477 // that is not the one being emptied keeps the copy it has there.
1478 self.detach(id, p);
1479 if vectors.get(id, &mut buf) {
1480 xs.extend_from_slice(&self.quant.rotate(&buf));
1481 kept.push(Member { id, tag });
1482 }
1483 }
1484 (kept, xs)
1485 }
1486
1487 /// The `n` partitions whose centroids are nearest `x`, nearest first.
1488 ///
1489 /// This measures against every centroid in the collection and it stays that
1490 /// way. Putting it through [`crate::coarse`] was tried and the recall it
1491 /// costs is not worth the time it saves, which the module doc there sets out
1492 /// with the numbers.
1493 /// The `n` partitions near `x`, as far as the coarse layer can tell.
1494 ///
1495 /// This is only for the sweep, and the reason it is allowed there and not on
1496 /// the search path is that it picks which partitions to look in rather than
1497 /// what the answer is. Every member the sweep then looks at is compared
1498 /// exactly against the centroids that just changed, so a neighbour the layer
1499 /// missed costs a few members not moving yet rather than a member moving to
1500 /// the wrong place, and the next split in that neighbourhood picks them up.
1501 /// The distinction matters because an approximate decision inside the sweep
1502 /// is the one thing [`crate::coarse`] says out loud must not happen.
1503 fn roughly_near_partitions(&self, x: &[f32], n: usize) -> Vec<usize> {
1504 if !self.coarse.ready() {
1505 return self.near_partitions(x, n);
1506 }
1507 let mut short = Vec::new();
1508 self.coarse.shortlist(x, self.dim(), &mut short);
1509 let mut by: Vec<(usize, f32)> = short
1510 .iter()
1511 .map(|&p| (p as usize, sqdist(x, self.centroid(p as usize))))
1512 .collect();
1513 let n = n.min(by.len());
1514 by.select_nth_unstable_by(n.saturating_sub(1), |a, b| a.1.total_cmp(&b.1));
1515 by.truncate(n);
1516 by.sort_by(|a, b| a.1.total_cmp(&b.1));
1517 by.into_iter().map(|(p, _)| p).collect()
1518 }
1519
1520 fn near_partitions(&self, x: &[f32], n: usize) -> Vec<usize> {
1521 let mut by: Vec<(usize, f32)> = (0..self.postings.len())
1522 .map(|p| (p, sqdist(x, self.centroid(p))))
1523 .collect();
1524 let n = n.min(by.len());
1525 by.select_nth_unstable_by(n.saturating_sub(1), |a, b| a.1.total_cmp(&b.1));
1526 by.truncate(n);
1527 by.sort_by(|a, b| a.1.total_cmp(&b.1));
1528 by.into_iter().map(|(p, _)| p).collect()
1529 }
1530
1531 /// The partitions in the order a search would probe them, nearest centroid
1532 /// first, all of them.
1533 ///
1534 /// Test only, and it exists for [`miss`](crate::miss), which asks how far
1535 /// down this order a query's true neighbours sit. That is the measurement
1536 /// that says whether the recall gate wants better partitions or a better
1537 /// estimator, and it cannot be taken from outside the crate because the
1538 /// probe order is not something a caller has any business seeing.
1539 #[cfg(test)]
1540 pub(crate) fn probe_order(&self, q: &[f32], into: &mut Vec<usize>) {
1541 let u = self.quant.rotate(q);
1542 *into = self.near_partitions(&u, self.postings.len());
1543 }
1544
1545 /// Which partition holds `id`, if any.
1546 #[cfg(test)]
1547 pub(crate) fn holder(&self, id: u64) -> Option<usize> {
1548 self.any_place(id).map(|s| s.partition as usize)
1549 }
1550
1551 /// The partition `x` belongs to, as far as the coarse layer can tell.
1552 fn roughly_nearest(&self, x: &[f32], short: &mut Vec<u32>) -> usize {
1553 if !self.coarse.ready() {
1554 return self.nearest(x);
1555 }
1556 self.coarse.shortlist(x, self.dim(), short);
1557 short
1558 .iter()
1559 .map(|&p| (p as usize, sqdist(x, self.centroid(p as usize))))
1560 .min_by(|a, b| a.1.total_cmp(&b.1))
1561 .map_or(0, |(p, _)| p)
1562 }
1563
1564 /// The partition `x` belongs to.
1565 fn nearest(&self, x: &[f32]) -> usize {
1566 (0..self.postings.len())
1567 .map(|p| (p, sqdist(x, self.centroid(p))))
1568 .min_by(|a, b| a.1.total_cmp(&b.1))
1569 .map_or(0, |(p, _)| p)
1570 }
1571
1572 fn centroid(&self, p: usize) -> &[f32] {
1573 let dim = self.dim();
1574 &self.centroids[p * dim..(p + 1) * dim]
1575 }
1576
1577 /// A new empty partition around `centroid`, which is already rotated.
1578 fn add_partition(&mut self, centroid: &[f32]) -> usize {
1579 let dim = self.quant.dim();
1580 self.centroids.extend_from_slice(centroid);
1581 self.postings.push(Posting::default());
1582 let p = self.postings.len() - 1;
1583 self.coarse.added(p, centroid, dim);
1584 self.note(p);
1585 self.refresh_coarse();
1586 p
1587 }
1588
1589 /// Rebuild the coarse layer if the partition count has moved far enough
1590 /// since the anchors were last chosen.
1591 fn refresh_coarse(&mut self) {
1592 let n = self.postings.len();
1593 if self.coarse.stale(n) {
1594 let dim = self.quant.dim();
1595 self.coarse.rebuild(&self.centroids, dim, n);
1596 }
1597 }
1598
1599 /// Drop an empty partition, moving the last one into its place.
1600 fn drop_partition(&mut self, p: usize) {
1601 debug_assert_eq!(self.postings[p].len(), 0, "a partition is emptied first");
1602 let dim = self.dim();
1603 let last = self.postings.len() - 1;
1604 self.coarse.dropped(p);
1605 self.postings.swap_remove(p);
1606 for i in 0..dim {
1607 self.centroids[p * dim + i] = self.centroids[last * dim + i];
1608 }
1609 self.centroids.truncate(last * dim);
1610 if p != last {
1611 // The partition that used to be last is at `p` now, so everything
1612 // filed under it has to be told. It is the copy in `last` that
1613 // moves, not the member, so a replicated id keeps its other copies
1614 // pointing where they already point.
1615 for i in 0..self.postings[p].len() {
1616 let id = self.postings[p].ids[i];
1617 if let Some(at) = self.placed_at(id, last) {
1618 self.places[at as usize].partition = p as u32;
1619 }
1620 }
1621 self.note(p);
1622 }
1623 self.refresh_coarse();
1624 }
1625
1626 /// Append a member to a partition. `x` is rotated.
1627 ///
1628 /// A partition that already holds a copy of `id` keeps the one it has, so
1629 /// that the two maintenance paths that can hand the same member to the same
1630 /// partition twice, a merge into a partition the member was replicated into
1631 /// and a sweep that moves it there, cannot produce a posting with the same
1632 /// id in it twice.
1633 fn place(&mut self, p: usize, id: u64, tag: u64, x: &[f32]) {
1634 let dim = self.dim();
1635 let width = self.quant.code_bytes();
1636 let slot = self.postings[p].len();
1637 if !self.attach(id, p, slot) {
1638 return;
1639 }
1640 self.postings[p].codes.resize((slot + 1) * width, 0);
1641 let centroid = &self.centroids[p * dim..(p + 1) * dim];
1642 let coded = self.quant.encode_rotated(
1643 x,
1644 centroid,
1645 &mut self.postings[p].codes[slot * width..(slot + 1) * width],
1646 );
1647 self.postings[p].ids.push(id);
1648 self.postings[p].tags.push(tag);
1649 self.postings[p].meta.push(coded);
1650 self.note(p);
1651 }
1652
1653 // -- the placement chain -------------------------------------------------
1654 //
1655 // Every site that used to write `self.at` goes through one of these, because
1656 // with replication the question is almost never about an id. It is about one
1657 // copy of an id, the one in a particular partition, and the difference only
1658 // shows up as a corrupt index a long way from where it was caused.
1659
1660 /// Record that `id` has a copy at `(p, slot)`, saying whether it is new.
1661 ///
1662 /// A partition already holding a copy is left alone rather than given a
1663 /// second one. Nothing on the insert path asks for that, but the maintenance
1664 /// paths can: a member replicated into two partitions that are then merged
1665 /// into each other would otherwise arrive twice, and a duplicate inside one
1666 /// posting is the one shape the rest of this cannot cope with, because a
1667 /// delete would take out one copy and leave the other.
1668 fn attach(&mut self, id: u64, p: usize, slot: usize) -> bool {
1669 let head = self.at.get(&id).copied().unwrap_or(END);
1670 let mut walk = head;
1671 while walk != END {
1672 if self.places[walk as usize].partition as usize == p {
1673 return false;
1674 }
1675 walk = self.places[walk as usize].next;
1676 }
1677 let place = Place {
1678 partition: p as u32,
1679 slot: slot as u32,
1680 next: head,
1681 };
1682 let at = if self.free == END {
1683 self.places.push(place);
1684 (self.places.len() - 1) as u32
1685 } else {
1686 let at = self.free;
1687 self.free = self.places[at as usize].next;
1688 self.places[at as usize] = place;
1689 at
1690 };
1691 self.at.insert(id, at);
1692 true
1693 }
1694
1695 /// Forget the copy of `id` in partition `p`, saying whether there was one.
1696 fn detach(&mut self, id: u64, p: usize) -> bool {
1697 let Some(&head) = self.at.get(&id) else {
1698 return false;
1699 };
1700 let mut prev = END;
1701 let mut walk = head;
1702 while walk != END {
1703 let this = self.places[walk as usize];
1704 if this.partition as usize == p {
1705 if prev == END {
1706 if this.next == END {
1707 self.at.remove(&id);
1708 } else {
1709 self.at.insert(id, this.next);
1710 }
1711 } else {
1712 self.places[prev as usize].next = this.next;
1713 }
1714 self.places[walk as usize].next = self.free;
1715 self.free = walk;
1716 return true;
1717 }
1718 prev = walk;
1719 walk = this.next;
1720 }
1721 false
1722 }
1723
1724 /// Forget every copy of `id`, saying whether there were any.
1725 fn detach_all(&mut self, id: u64) -> bool {
1726 let Some(head) = self.at.remove(&id) else {
1727 return false;
1728 };
1729 let mut walk = head;
1730 while walk != END {
1731 let next = self.places[walk as usize].next;
1732 self.places[walk as usize].next = self.free;
1733 self.free = walk;
1734 walk = next;
1735 }
1736 true
1737 }
1738
1739 /// Where the copy of `id` in partition `p` is, if there is one.
1740 fn placed_at(&self, id: u64, p: usize) -> Option<u32> {
1741 let mut walk = self.at.get(&id).copied().unwrap_or(END);
1742 while walk != END {
1743 if self.places[walk as usize].partition as usize == p {
1744 return Some(walk);
1745 }
1746 walk = self.places[walk as usize].next;
1747 }
1748 None
1749 }
1750
1751 /// Say that the copy of `id` in partition `p` is at slot `s` now, which is
1752 /// what a pull leaves behind when it moves the last member into a hole.
1753 fn reslot(&mut self, id: u64, p: usize, s: usize) {
1754 if let Some(at) = self.placed_at(id, p) {
1755 self.places[at as usize].slot = s as u32;
1756 } else {
1757 debug_assert!(
1758 false,
1759 "id {id} is in partition {p} and the map does not say so"
1760 );
1761 }
1762 }
1763
1764 /// How many partitions hold a copy of `id`.
1765 #[cfg(test)]
1766 fn placements_of(&self, id: u64) -> usize {
1767 let mut walk = self.at.get(&id).copied().unwrap_or(END);
1768 let mut n = 0;
1769 while walk != END {
1770 n += 1;
1771 walk = self.places[walk as usize].next;
1772 }
1773 n
1774 }
1775
1776 /// Any one copy of `id`, for the questions that do not care which.
1777 fn any_place(&self, id: u64) -> Option<Place> {
1778 self.at.get(&id).map(|&at| self.places[at as usize])
1779 }
1780
1781 /// Every copy of `id`, collected because the callers that want them all are
1782 /// about to borrow the index mutably.
1783 fn every_place(&self, id: u64, into: &mut Vec<Place>) {
1784 into.clear();
1785 let mut walk = self.at.get(&id).copied().unwrap_or(END);
1786 while walk != END {
1787 let place = self.places[walk as usize];
1788 into.push(place);
1789 walk = place.next;
1790 }
1791 }
1792
1793 /// Take slot `s` out of partition `p`, returning the id that moved into it.
1794 fn pull(&mut self, p: usize, s: usize) -> Option<u64> {
1795 let width = self.quant.code_bytes();
1796 let posting = &mut self.postings[p];
1797 let last = posting.len() - 1;
1798 posting.ids.swap_remove(s);
1799 posting.tags.swap_remove(s);
1800 posting.meta.swap_remove(s);
1801 if s != last {
1802 let (head, tail) = posting.codes.split_at_mut(last * width);
1803 head[s * width..(s + 1) * width].copy_from_slice(&tail[..width]);
1804 }
1805 posting.codes.truncate(last * width);
1806 let moved = (s != last).then(|| posting.ids[s]);
1807 self.note(p);
1808 moved
1809 }
1810
1811 /// The same, keeping the map straight, for the paths that are about to put
1812 /// the member somewhere else.
1813 fn pull_and_forget(&mut self, p: usize, s: usize) {
1814 let id = self.postings[p].ids[s];
1815 self.detach(id, p);
1816 if let Some(moved) = self.pull(p, s) {
1817 self.reslot(moved, p, s);
1818 }
1819 }
1820}
1821
1822/// A member on its way from one partition to another, which is the only time
1823/// its id and its tag travel together without a posting around them.
1824#[derive(Clone, Copy)]
1825struct Member {
1826 id: u64,
1827 tag: u64,
1828}
1829
1830#[derive(Debug, PartialEq, Eq)]
1831enum Job {
1832 Split(usize),
1833 Merge(usize),
1834}
1835
1836/// Two means over a set of vectors laid out end to end.
1837///
1838/// The seeds are the member furthest from the middle and then the member
1839/// furthest from that one, which is deterministic, needs no generator, and
1840/// starts on the axis the cloud is actually longest along. Eight rounds is
1841/// past where this stops moving on anything shaped like an embedding.
1842fn two_means(xs: &[f32], dim: usize) -> (Vec<f32>, Vec<f32>) {
1843 let n = xs.len() / dim;
1844 let mut middle = vec![0.0f32; dim];
1845 for i in 0..n {
1846 for (m, c) in middle.iter_mut().zip(&xs[i * dim..(i + 1) * dim]) {
1847 *m += c;
1848 }
1849 }
1850 for m in &mut middle {
1851 *m /= n as f32;
1852 }
1853 let far = |from: &[f32]| {
1854 (0..n)
1855 .max_by(|&i, &j| {
1856 sqdist(from, &xs[i * dim..(i + 1) * dim])
1857 .total_cmp(&sqdist(from, &xs[j * dim..(j + 1) * dim]))
1858 })
1859 .unwrap_or(0)
1860 };
1861 let i = far(&middle);
1862 let mut a = xs[i * dim..(i + 1) * dim].to_vec();
1863 let j = far(&a);
1864 let mut b = xs[j * dim..(j + 1) * dim].to_vec();
1865
1866 for _ in 0..8 {
1867 let mut sums = (vec![0.0f32; dim], vec![0.0f32; dim]);
1868 let mut counts = (0usize, 0usize);
1869 for i in 0..n {
1870 let x = &xs[i * dim..(i + 1) * dim];
1871 if sqdist(x, &a) <= sqdist(x, &b) {
1872 for (s, c) in sums.0.iter_mut().zip(x) {
1873 *s += c;
1874 }
1875 counts.0 += 1;
1876 } else {
1877 for (s, c) in sums.1.iter_mut().zip(x) {
1878 *s += c;
1879 }
1880 counts.1 += 1;
1881 }
1882 }
1883 // A side that ended up with nothing keeps the seed it had, because a
1884 // mean of no points is not a place and the next round would put every
1885 // member on the other side for ever.
1886 if counts.0 > 0 {
1887 for (m, s) in a.iter_mut().zip(&sums.0) {
1888 *m = s / counts.0 as f32;
1889 }
1890 }
1891 if counts.1 > 0 {
1892 for (m, s) in b.iter_mut().zip(&sums.1) {
1893 *m = s / counts.1 as f32;
1894 }
1895 }
1896 }
1897 (a, b)
1898}
1899
1900/// One candidate, ordered by its estimated distance.
1901///
1902/// The tie break on the id is not decoration. Two members of the same partition
1903/// can get the same estimate out of codes that are 16 bytes wide, and without a
1904/// tie break which of them survives depends on the order the heap happened to
1905/// be in, which makes a search answer depend on the insertion history of the
1906/// collection rather than on the collection.
1907#[derive(PartialEq)]
1908struct Ranked {
1909 at: f32,
1910 id: u64,
1911}
1912
1913impl Eq for Ranked {}
1914
1915impl Ord for Ranked {
1916 fn cmp(&self, other: &Ranked) -> std::cmp::Ordering {
1917 self.at.total_cmp(&other.at).then(self.id.cmp(&other.id))
1918 }
1919}
1920
1921impl PartialOrd for Ranked {
1922 fn partial_cmp(&self, other: &Ranked) -> Option<std::cmp::Ordering> {
1923 Some(self.cmp(other))
1924 }
1925}
1926
1927/// The best `want` candidates seen so far, and nothing else.
1928///
1929/// The scan used to push every member of every partition it read into one
1930/// vector and then select from it, which at probe 64 is 24 thousand entries
1931/// pushed and 24 thousand selected over to keep 160. That is 384 kilobytes of
1932/// writes per search and it was a tenth of the search's time.
1933///
1934/// A bounded heap makes the common case one comparison. Once `want` candidates
1935/// are in, a member is only touched further if it beats the worst of them,
1936/// which after the first partition or two is a small fraction of them, and a
1937/// member that loses never has its id or its tag read at all.
1938struct Bounded {
1939 want: usize,
1940 heap: std::collections::BinaryHeap<Ranked>,
1941}
1942
1943impl Bounded {
1944 fn new(want: usize) -> Bounded {
1945 Bounded {
1946 want,
1947 heap: std::collections::BinaryHeap::with_capacity(want + 1),
1948 }
1949 }
1950
1951 /// Whether there are already `want` answers, which is what says a search
1952 /// that was widening for a filter can stop widening.
1953 fn full(&self) -> bool {
1954 self.heap.len() >= self.want
1955 }
1956
1957 /// Whether `at` could still be one of the answers.
1958 #[inline]
1959 fn wants(&self, at: f32) -> bool {
1960 match self.heap.peek() {
1961 Some(worst) if self.heap.len() >= self.want => at < worst.at,
1962 _ => true,
1963 }
1964 }
1965
1966 fn put(&mut self, id: u64, at: f32) {
1967 if self.heap.len() >= self.want {
1968 self.heap.pop();
1969 }
1970 self.heap.push(Ranked { at, id });
1971 }
1972
1973 /// The answers, nearest first.
1974 fn sorted(self) -> Vec<(u64, f32)> {
1975 self.heap
1976 .into_sorted_vec()
1977 .into_iter()
1978 .map(|r| (r.id, r.at))
1979 .collect()
1980 }
1981}
1982
1983#[cfg(test)]
1984mod tests {
1985 use super::*;
1986 use yo_common::Rng;
1987
1988 /// What `job` used to be: two passes over every partition. The candidate
1989 /// lists are only worth having if they give the same answer, so the slow
1990 /// version stays here as the thing the fast one is checked against.
1991 fn slow_job(ix: &Partitions) -> Option<Job> {
1992 let big = (0..ix.postings.len())
1993 .filter(|&p| ix.postings[p].len() > ix.postings[p].stuck)
1994 .max_by_key(|&p| ix.postings[p].len());
1995 if let Some(big) = big
1996 && ix.postings[big].len() > ix.tuning.posting * 2
1997 {
1998 return Some(Job::Split(big));
1999 }
2000 if ix.postings.len() > 1 {
2001 let small = (0..ix.postings.len()).min_by_key(|&p| ix.postings[p].len())?;
2002 if ix.postings[small].len() * 4 < ix.tuning.posting {
2003 return Some(Job::Merge(small));
2004 }
2005 }
2006 None
2007 }
2008
2009 /// Every partition that qualifies has to be on a list, or maintenance stops
2010 /// happening and the index quietly rots. Deletes are in here on purpose,
2011 /// because a delete is the only thing that pushes a partition down through
2012 /// the merge threshold and it is also what makes a partition disappear and
2013 /// renumber the one that was last.
2014 #[test]
2015 fn the_candidate_lists_answer_what_the_two_passes_answered() {
2016 let (n, dim, posting) = shrunk(3000, 16, Tuning::default().posting);
2017 let store = corpus(dim, n, 12, 0x105E);
2018 let tuning = Tuning {
2019 posting,
2020 ..Tuning::default()
2021 };
2022 let mut ix = Partitions::new(dim, Bits::One, 7, tuning);
2023 let mut rng = Rng::new(0x105F);
2024 let mut live: Vec<u64> = Vec::new();
2025 for id in 0..n as u64 {
2026 ix.insert(id, &store.0[id as usize]);
2027 live.push(id);
2028 if id % 7 == 3 && !live.is_empty() {
2029 let at = rng.below(live.len());
2030 let gone = live.swap_remove(at);
2031 ix.remove(gone);
2032 }
2033 // Once before maintenance runs and once after, because the lists
2034 // are written by both and the state in between is the one a stale
2035 // entry would survive in.
2036 assert_eq!(
2037 ix.job(),
2038 slow_job(&ix),
2039 "after {id} inserts, before maintaining"
2040 );
2041 ix.maintain(&store, 8);
2042 assert_eq!(
2043 ix.job(),
2044 slow_job(&ix),
2045 "after {id} inserts, after maintaining"
2046 );
2047 assert_eq!(ix.needs_maintenance(), slow_job(&ix).is_some());
2048 }
2049 assert!(ix.postings.len() > 5, "the test never split anything");
2050 }
2051
2052 /// The record log, for a test: every vector by id, where the id is where it
2053 /// sits.
2054 struct Store(Vec<Vec<f32>>);
2055
2056 impl Vectors for Store {
2057 fn get(&self, id: u64, into: &mut [f32]) -> bool {
2058 match self.0.get(id as usize) {
2059 Some(v) => {
2060 into.copy_from_slice(v);
2061 true
2062 }
2063 None => false,
2064 }
2065 }
2066 }
2067
2068 /// A store with a hole in it, for the case where the log forgot something
2069 /// the index still thinks it has.
2070 struct Holey(Vec<Vec<f32>>, u64);
2071
2072 impl Vectors for Holey {
2073 fn get(&self, id: u64, into: &mut [f32]) -> bool {
2074 if id == self.1 {
2075 return false;
2076 }
2077 match self.0.get(id as usize) {
2078 Some(v) => {
2079 into.copy_from_slice(v);
2080 true
2081 }
2082 None => false,
2083 }
2084 }
2085 }
2086
2087 /// Vectors with the two things real embeddings have and uniform noise does
2088 /// not: a few coordinates carrying most of the energy, so that two vectors
2089 /// are genuinely near each other rather than all being equally far apart,
2090 /// and clusters, so that the partitioning has something to get right.
2091 ///
2092 /// A corpus without the first of those is not a hard test, it is an
2093 /// impossible one. The nearest ten of three hundred uniform points are
2094 /// arbitrary, no quantiser can pick them out, and the recall it measures
2095 /// says nothing about the index.
2096 fn corpus(dim: usize, n: usize, clusters: usize, seed: u64) -> Store {
2097 let mut rng = Rng::new(seed);
2098 let centres: Vec<Vec<f32>> = (0..clusters).map(|_| draw(dim, &mut rng)).collect();
2099 Store(
2100 (0..n)
2101 .map(|i| {
2102 let off = draw(dim, &mut rng);
2103 let mut v: Vec<f32> = centres[i % clusters]
2104 .iter()
2105 .zip(&off)
2106 .map(|(c, o)| c + o * 0.7)
2107 .collect();
2108 unit(&mut v);
2109 v
2110 })
2111 .collect(),
2112 )
2113 }
2114
2115 /// A corpus size, a width and a posting size, shrunk together for Miri.
2116 ///
2117 /// The corpus and the posting have to move together. What the tests using
2118 /// this are about is the shape the index takes: how many partitions there
2119 /// are, that a posting splits when it grows and merges when it shrinks, that
2120 /// a member near a boundary ends up copied into a second partition. All of
2121 /// that is set by the ratio of those two numbers rather than by either of
2122 /// them, so cutting both by twenty leaves every one of those claims where it
2123 /// was and takes a twentieth of the arithmetic to get there.
2124 ///
2125 /// Cutting only the corpus is the obvious half of this and it is the wrong
2126 /// half. It gives a collection that never splits and a set of tests that
2127 /// pass without having looked at anything, which is worse than leaving them
2128 /// out. That mistake is what kept `yo-index` out of the Miri run for a
2129 /// month, from the other direction: a count that had been shrunk sitting
2130 /// next to a count that had not.
2131 ///
2132 /// The width is the third one because it is the one that costs the most and
2133 /// says the least. A rotation is `dim` squared multiplications and every
2134 /// insert and every query pays for one, so a corpus in 96 dimensions is a
2135 /// hundred and forty times the arithmetic of the same corpus in 8 for a
2136 /// claim that is about placement rather than about geometry. Eight is where
2137 /// the floor is because `corpus` now puts the energy in at least one
2138 /// coordinate at any width, and a corpus with no heavy coordinate at all is
2139 /// not a smaller version of the test, it is uniform noise where nothing is
2140 /// near anything.
2141 ///
2142 /// The floors on the other two are there because below a hundred or so
2143 /// members, or six or so to a posting, a split stops being a thing that
2144 /// happens inside the index and becomes the whole index. Where a floor
2145 /// bites, the ratio moves a little and the partition count comes out higher
2146 /// rather than lower, which is the safe direction: more partitions is more
2147 /// of the thing being tested.
2148 ///
2149 /// Tests whose claim is a number rather than a shape do not come through
2150 /// here. A recall figure, a saving measured over a hundred queries, a
2151 /// drifted-member count against another drifted-member count: shrinking any
2152 /// of those leaves a test that still passes and no longer means anything, so
2153 /// those are skipped under Miri and each one says so.
2154 fn shrunk(n: usize, dim: usize, posting: usize) -> (usize, usize, usize) {
2155 if cfg!(miri) {
2156 ((n / 20).max(120), dim.min(8), (posting / 20).max(6))
2157 } else {
2158 (n, dim, posting)
2159 }
2160 }
2161
2162 /// One vector of the shape above, unit length.
2163 fn draw(dim: usize, rng: &mut Rng) -> Vec<f32> {
2164 let mut v: Vec<f32> = (0..dim)
2165 .map(|i| {
2166 let u = (rng.next_u64() >> 40) as f32 / (1u32 << 24) as f32;
2167 // At least one, so that a narrow corpus is a smaller version of
2168 // a wide one rather than uniform noise where nothing is near
2169 // anything and no quantiser can pick out a nearest anything.
2170 let heavy = if i < (dim / 16).max(1) { 6.0 } else { 1.0 };
2171 (u * 2.0 - 1.0) * heavy
2172 })
2173 .collect();
2174 unit(&mut v);
2175 v
2176 }
2177
2178 fn unit(v: &mut [f32]) {
2179 let len = v.iter().map(|c| c * c).sum::<f32>().sqrt();
2180 for c in v {
2181 *c /= len;
2182 }
2183 }
2184
2185 fn truth(store: &Store, q: &[f32], k: usize) -> Vec<u64> {
2186 let mut all: Vec<(u64, f32)> = store
2187 .0
2188 .iter()
2189 .enumerate()
2190 .map(|(i, v)| (i as u64, sqdist(q, v)))
2191 .collect();
2192 all.sort_by(|a, b| a.1.total_cmp(&b.1));
2193 all.truncate(k);
2194 all.into_iter().map(|(i, _)| i).collect()
2195 }
2196
2197 /// Build an index over a whole store, running maintenance as it goes the way
2198 /// a maintenance slice would.
2199 fn build(store: &Store, dim: usize, tuning: Tuning) -> Partitions {
2200 let mut ix = Partitions::new(dim, Bits::One, 7, tuning);
2201 for (i, v) in store.0.iter().enumerate() {
2202 ix.insert(i as u64, v);
2203 if i % 64 == 0 {
2204 ix.maintain(store, 4096);
2205 }
2206 }
2207 ix.maintain(store, 1 << 20);
2208 ix
2209 }
2210
2211 /// How often the true `k` nearest come back.
2212 fn recall(ix: &Partitions, store: &Store, k: usize, queries: usize) -> f32 {
2213 let mut hits = 0usize;
2214 for i in 0..queries {
2215 // A query near a real vector rather than anywhere at all, because
2216 // that is what a search looks like.
2217 let q = &store.0[i * 7 % store.0.len()];
2218 let want = truth(store, q, k);
2219 let got: Vec<u64> = ix.search(q, k, store).into_iter().map(|h| h.id).collect();
2220 hits += want.iter().filter(|id| got.contains(id)).count();
2221 }
2222 hits as f32 / (queries * k) as f32
2223 }
2224
2225 /// Everything the index believes about itself, checked.
2226 fn consistent(ix: &Partitions) {
2227 assert_eq!(ix.centroids.len(), ix.postings.len() * ix.dim());
2228 let width = ix.quant.code_bytes();
2229 let mut seen = 0usize;
2230 for (p, posting) in ix.postings.iter().enumerate() {
2231 assert_eq!(posting.codes.len(), posting.len() * width, "partition {p}");
2232 assert_eq!(posting.meta.len(), posting.len(), "partition {p}");
2233 assert_eq!(posting.tags.len(), posting.len(), "partition {p}");
2234 let mut here = HashSet::new();
2235 for (s, id) in posting.ids.iter().enumerate() {
2236 assert!(here.insert(*id), "id {id} is twice in partition {p}");
2237 let at = ix
2238 .placed_at(*id, p)
2239 .expect("every member is in the map, under the partition holding it");
2240 assert_eq!(ix.places[at as usize].slot as usize, s, "id {id}");
2241 seen += 1;
2242 }
2243 }
2244 // Every placement points at a member, as many placements as there are
2245 // members, and the free list accounts for the rest of the arena. A
2246 // replicated id makes the first of those the interesting one: a chain
2247 // that kept an entry for a copy that was pulled would still look right
2248 // from the posting's side, and the count is what catches it.
2249 let mut held = 0usize;
2250 for (&id, &head) in &ix.at {
2251 let mut walk = head;
2252 let mut mine = HashSet::new();
2253 while walk != END {
2254 let place = ix.places[walk as usize];
2255 let p = place.partition as usize;
2256 assert!(mine.insert(p), "id {id} is filed twice under partition {p}");
2257 assert!(
2258 p < ix.postings.len(),
2259 "id {id} is filed under partition {p}"
2260 );
2261 assert_eq!(
2262 ix.postings[p].ids[place.slot as usize], id,
2263 "id {id} is filed at a slot holding something else"
2264 );
2265 held += 1;
2266 walk = place.next;
2267 }
2268 }
2269 assert_eq!(seen, held, "the map and the postings disagree on the count");
2270 let mut spare = 0usize;
2271 let mut walk = ix.free;
2272 while walk != END {
2273 spare += 1;
2274 assert!(spare <= ix.places.len(), "the free list has a cycle in it");
2275 walk = ix.places[walk as usize].next;
2276 }
2277 assert_eq!(held + spare, ix.places.len(), "the arena has leaked");
2278 }
2279
2280 /// Build an index where every vector carries a tag, so the filter has
2281 /// something to meet.
2282 fn build_tagged(
2283 store: &Store,
2284 dim: usize,
2285 tuning: Tuning,
2286 tag: impl Fn(u64) -> u64,
2287 ) -> Partitions {
2288 let mut ix = Partitions::new(dim, Bits::One, 7, tuning);
2289 for (i, v) in store.0.iter().enumerate() {
2290 ix.insert_tagged(i as u64, v, tag(i as u64));
2291 if i % 64 == 0 {
2292 ix.maintain(store, 4096);
2293 }
2294 }
2295 ix.maintain(store, 1 << 20);
2296 ix
2297 }
2298
2299 /// The whole point of pushing a filter into the scan, measured against the
2300 /// thing it replaces.
2301 ///
2302 /// One document in fifty carries the tag. Filtering inside the scan finds
2303 /// the true ten of those. Taking the best forty by vector and then throwing
2304 /// away the ones that do not match, which is what a search that cannot push
2305 /// a filter down has to do, finds almost none of them, and the ones it
2306 /// misses were nearer than the ones it kept.
2307 #[test]
2308 #[cfg_attr(
2309 miri,
2310 ignore = "the count is the claim: one document in fifty of three thousand, and the whole point is what the near misses were, which needs a corpus with near misses in it"
2311 )]
2312 fn a_filter_in_the_scan_finds_what_a_filter_after_it_cannot() {
2313 let dim = 96;
2314 let store = corpus(dim, 3000, 12, 47);
2315 let tuning = Tuning {
2316 posting: 64,
2317 ..Tuning::default()
2318 };
2319 let wanted = |id: u64| id.is_multiple_of(50);
2320 let ix = build_tagged(&store, dim, tuning, |id| u64::from(wanted(id)));
2321
2322 let (mut pushed, mut after) = (0usize, 0usize);
2323 let k = 10;
2324 for i in 0..40 {
2325 let q = &store.0[i * 71 % store.0.len()];
2326
2327 // What the answer is: brute force over the members that match.
2328 let mut all: Vec<(u64, f32)> = store
2329 .0
2330 .iter()
2331 .enumerate()
2332 .filter(|(id, _)| wanted(*id as u64))
2333 .map(|(id, v)| (id as u64, sqdist(q, v)))
2334 .collect();
2335 all.sort_by(|a, b| a.1.total_cmp(&b.1));
2336 let want: Vec<u64> = all[..k].iter().map(|(id, _)| *id).collect();
2337
2338 let got: Vec<u64> = ix
2339 .search_where(q, k, &|tag: u64| tag == 1, &store)
2340 .into_iter()
2341 .map(|h| h.id)
2342 .collect();
2343 pushed += want.iter().filter(|id| got.contains(id)).count();
2344
2345 let late: Vec<u64> = ix
2346 .search(q, k * tuning.rerank, &store)
2347 .into_iter()
2348 .map(|h| h.id)
2349 .filter(|id| wanted(*id))
2350 .take(k)
2351 .collect();
2352 after += want.iter().filter(|id| late.contains(id)).count();
2353 }
2354 let (pushed, after) = (pushed as f32 / 400.0, after as f32 / 400.0);
2355 assert!(pushed >= 0.95, "pushing the filter down gave {pushed}");
2356 assert!(
2357 after < pushed / 2.0,
2358 "filtering afterwards gave {after} against {pushed}, which is not the point being made"
2359 );
2360 }
2361
2362 /// The second test decides, and the scan keeps widening until it has `k` of
2363 /// what the second test wants rather than `k` of what the tag wants.
2364 ///
2365 /// This is what a filter whose tag is only a summary looks like: the tag
2366 /// lets a superset through, so an answer that passes it and fails the exact
2367 /// test must not have cost an answer that passes both.
2368 #[test]
2369 #[cfg_attr(
2370 miri,
2371 ignore = "the count is the claim: recall at ten where the tag lets one in ten through and the exact test keeps one in fifty"
2372 )]
2373 fn the_exact_test_decides_and_the_scan_widens_for_it() {
2374 struct Summary;
2375
2376 impl Filter for Summary {
2377 fn allows(&self, tag: u64) -> bool {
2378 // One in ten, which is what a bit that several values landed on
2379 // looks like from inside the scan.
2380 tag == 1
2381 }
2382
2383 fn exact(&self, id: u64) -> bool {
2384 // One in fifty, and a subset of what the tag allowed, which is
2385 // the direction a summary is allowed to be wrong in.
2386 id.is_multiple_of(50)
2387 }
2388 }
2389
2390 let dim = 64;
2391 let store = corpus(dim, 3000, 9, 71);
2392 let tuning = Tuning {
2393 posting: 64,
2394 ..Tuning::default()
2395 };
2396 let ix = build_tagged(&store, dim, tuning, |id| u64::from(id.is_multiple_of(10)));
2397
2398 let k = 10;
2399 let mut found = 0usize;
2400 for i in 0..20 {
2401 let q = &store.0[i * 131 % store.0.len()];
2402 let mut all: Vec<(u64, f32)> = store
2403 .0
2404 .iter()
2405 .enumerate()
2406 .map(|(id, v)| (id as u64, sqdist(q, v)))
2407 .filter(|(id, _)| id.is_multiple_of(50))
2408 .collect();
2409 all.sort_by(|a, b| a.1.total_cmp(&b.1));
2410 let want: Vec<u64> = all[..k].iter().map(|(id, _)| *id).collect();
2411
2412 let got: Vec<u64> = ix
2413 .search_where(q, k, &Summary, &store)
2414 .into_iter()
2415 .map(|h| h.id)
2416 .collect();
2417 assert!(
2418 got.iter().all(|id| id.is_multiple_of(50)),
2419 "the exact test did not decide: {got:?}"
2420 );
2421 found += want.iter().filter(|id| got.contains(id)).count();
2422 }
2423 let recall = found as f32 / (20.0 * k as f32);
2424 assert!(recall >= 0.9, "two stage filtering gave {recall}");
2425 }
2426
2427 #[test]
2428 fn a_filter_that_matches_nothing_answers_nothing() {
2429 let (n, dim, posting) = shrunk(500, 64, Tuning::default().posting);
2430 let store = corpus(dim, n, 4, 53);
2431 let tuning = Tuning {
2432 posting,
2433 ..Tuning::default()
2434 };
2435 let ix = build_tagged(&store, dim, tuning, |_| 1);
2436 assert!(
2437 ix.search_where(&store.0[0], 10, &|tag: u64| tag == 2, &store)
2438 .is_empty()
2439 );
2440 // And the same filter matching everything is the unfiltered answer.
2441 let all = ix.search_where(&store.0[0], 10, &Any, &store);
2442 assert_eq!(all, ix.search(&store.0[0], 10, &store));
2443 }
2444
2445 #[test]
2446 fn a_tag_survives_a_split_and_a_merge() {
2447 let (n, dim, posting) = shrunk(800, 64, 24);
2448 let store = corpus(dim, n, 6, 59);
2449 let tuning = Tuning {
2450 posting,
2451 ..Tuning::default()
2452 };
2453 let mut ix = build_tagged(&store, dim, tuning, |id| id * 7 + 1);
2454 assert!(ix.partitions() > 4, "it never split");
2455 for id in 0..n as u64 {
2456 assert_eq!(ix.tag(id), Some(id * 7 + 1), "id {id} after the splits");
2457 }
2458
2459 // Now shrink it until partitions merge, and the survivors keep theirs.
2460 // Forty left, whatever the corpus was, because forty is what makes the
2461 // postings small enough to merge and the survivors are counted rather
2462 // than sampled.
2463 let left = n as u64 - 40;
2464 for id in 0..left {
2465 ix.remove(id);
2466 }
2467 ix.maintain(&store, 1 << 20);
2468 consistent(&ix);
2469 for id in left..n as u64 {
2470 assert_eq!(ix.tag(id), Some(id * 7 + 1), "id {id} after the merges");
2471 }
2472 assert_eq!(ix.tag(0), None);
2473 }
2474
2475 /// A selective filter means the answers are not in the nearest partitions,
2476 /// and a search that will not look further returns fewer than it should.
2477 #[test]
2478 #[cfg_attr(
2479 miri,
2480 ignore = "the count is the claim: one member in a hundred of two thousand is twenty answers, and the assertion is that exact number"
2481 )]
2482 fn a_selective_filter_makes_the_search_look_further() {
2483 let dim = 64;
2484 let store = corpus(dim, 2000, 10, 61);
2485 let tuning = Tuning {
2486 posting: 32,
2487 ..Tuning::default()
2488 };
2489 let tag = |id: u64| u64::from(id.is_multiple_of(100));
2490 let ix = build_tagged(&store, dim, tuning, tag);
2491 let narrow = build_tagged(&store, dim, Tuning { widen: 1, ..tuning }, tag);
2492
2493 let mut wide_found = 0usize;
2494 let mut narrow_found = 0usize;
2495 for i in 0..20 {
2496 let q = &store.0[i * 91 % store.0.len()];
2497 wide_found += ix.search_where(q, 10, &|t: u64| t == 1, &store).len();
2498 narrow_found += narrow.search_where(q, 10, &|t: u64| t == 1, &store).len();
2499 }
2500 assert_eq!(
2501 wide_found, 200,
2502 "one in a hundred of two thousand is twenty"
2503 );
2504 assert!(
2505 narrow_found < wide_found,
2506 "not widening found {narrow_found} of {wide_found}"
2507 );
2508 }
2509
2510 #[test]
2511 fn a_signature_never_rejects_something_it_should_have_matched() {
2512 let english = Signature::of(&[("lang", b"en")]);
2513 let doc = Signature::of(&[("lang", b"en"), ("topic", b"finance"), ("year", b"2026")]);
2514 assert!(doc.covers(english));
2515 assert!(english.allows(doc.bits()));
2516 assert_eq!(Signature::from_bits(doc.bits()), doc);
2517
2518 // And over a lot of values, nothing that matches is ever turned away.
2519 for i in 0..500u32 {
2520 let value = i.to_string();
2521 let one = Signature::of(&[("id", value.as_bytes())]);
2522 let with = Signature::of(&[("id", value.as_bytes()), ("kind", b"page")]);
2523 assert!(with.covers(one), "value {value}");
2524 }
2525 }
2526
2527 #[test]
2528 fn an_empty_index_answers_nothing() {
2529 let ix = Partitions::new(32, Bits::One, 1, Tuning::default());
2530 let store = Store(Vec::new());
2531 assert!(ix.is_empty());
2532 assert_eq!(ix.partitions(), 0);
2533 assert!(ix.search(&[0.0; 32], 10, &store).is_empty());
2534 assert!(!ix.needs_maintenance());
2535 }
2536
2537 #[test]
2538 fn the_first_vector_is_the_first_partition() {
2539 let store = corpus(32, 1, 1, 3);
2540 let mut ix = Partitions::new(32, Bits::One, 1, Tuning::default());
2541 ix.insert(0, &store.0[0]);
2542 assert_eq!(ix.partitions(), 1);
2543 assert_eq!(ix.len(), 1);
2544 let hits = ix.search(&store.0[0], 5, &store);
2545 assert_eq!(hits.len(), 1);
2546 assert_eq!(hits[0].id, 0);
2547 assert!(hits[0].distance < 1e-6, "{}", hits[0].distance);
2548 consistent(&ix);
2549 }
2550
2551 #[test]
2552 #[cfg_attr(
2553 miri,
2554 ignore = "the count is the claim: recall at ten over two thousand vectors, and a recall figure over a corpus small enough to interpret is a number about nothing"
2555 )]
2556 fn a_search_finds_what_brute_force_finds() {
2557 let dim = 128;
2558 let store = corpus(dim, 2000, 12, 5);
2559 let ix = build(&store, dim, Tuning::default());
2560 assert!(ix.partitions() > 1, "it never split");
2561 consistent(&ix);
2562 let r = recall(&ix, &store, 10, 50);
2563 assert!(r >= 0.95, "recall at 10 was {r}");
2564 }
2565
2566 #[test]
2567 fn a_posting_that_grows_too_big_splits() {
2568 let (n, dim, posting) = shrunk(600, 64, 32);
2569 let tuning = Tuning {
2570 posting,
2571 ..Tuning::default()
2572 };
2573 let store = corpus(dim, n, 6, 9);
2574 let ix = build(&store, dim, tuning);
2575 assert!(
2576 ix.partitions() >= n / (posting * 2),
2577 "{n} vectors in {} partitions",
2578 ix.partitions()
2579 );
2580 // Over the limit is allowed for a posting that gave up, which is a
2581 // posting with no cut in it worth taking, and `stuck` is how it says so.
2582 // That case does not come up at six hundred vectors in sixty four
2583 // dimensions and does come up at a size Miri can afford, where the merge
2584 // threshold is low enough that a lopsided cut would be undone as fast as
2585 // it was made.
2586 for held in &ix.postings {
2587 assert!(
2588 held.len() <= posting * 2 || held.len() <= held.stuck,
2589 "a posting is {} long and did not give up splitting",
2590 held.len()
2591 );
2592 }
2593 consistent(&ix);
2594 }
2595
2596 #[test]
2597 fn a_posting_that_shrinks_merges() {
2598 let (n, dim, posting) = shrunk(600, 64, 32);
2599 let tuning = Tuning {
2600 posting,
2601 ..Tuning::default()
2602 };
2603 let store = corpus(dim, n, 6, 9);
2604 let mut ix = build(&store, dim, tuning);
2605 let grown = ix.partitions();
2606 assert!(grown > 4);
2607
2608 // Take away almost everything and let maintenance settle. Thirty left,
2609 // whatever the corpus was, because thirty is under the merge threshold
2610 // for any posting size this runs at.
2611 let left = n as u64 - 30;
2612 for id in 0..left {
2613 assert!(ix.remove(id));
2614 }
2615 ix.maintain(&store, 1 << 20);
2616 consistent(&ix);
2617 assert_eq!(ix.len(), 30);
2618 assert!(
2619 ix.partitions() < grown,
2620 "{} partitions for 30 vectors, was {grown}",
2621 ix.partitions()
2622 );
2623 // And it still answers.
2624 let last = n as u64 - 1;
2625 let hits = ix.search(&store.0[last as usize], 1, &store);
2626 assert_eq!(hits[0].id, last);
2627 }
2628
2629 /// Maintenance finishes, rather than taking turns with itself forever.
2630 ///
2631 /// A settled collection that nobody is writing to should have nothing left
2632 /// for `maintain` to do, and for a long time one shape of collection had an
2633 /// endless amount. A cloud with one outlier in it splits into the cloud and
2634 /// the outlier, the outlier is a partition of one so `merge` gives it back
2635 /// to its nearest centroid, which is the cloud it just came out of, and the
2636 /// cloud is over the limit again. Nothing about that is visible from inside
2637 /// either job. Both of them do exactly what they are for, the collection
2638 /// does not change, and every call to `maintain` spends its whole budget.
2639 ///
2640 /// It was found under Miri, where a hundred and sixty vectors took longer
2641 /// than fifteen minutes, and it is checked here across a spread of posting
2642 /// sizes because whether a given corpus falls into it depends on where the
2643 /// merge threshold lands relative to the cut two means happens to make.
2644 /// Twelve is the size that caught it. Six and sixteen, on the same vectors,
2645 /// settle in a few hundred.
2646 #[test]
2647 fn maintenance_runs_out_of_things_to_do() {
2648 for posting in [6, 8, 12, 16, 24] {
2649 let store = corpus(16, 160, 4, 11);
2650 let tuning = Tuning {
2651 posting,
2652 ..Tuning::default()
2653 };
2654 let mut ix = build(&store, 16, tuning);
2655 let left = ix.maintain(&store, 1 << 20);
2656 assert_eq!(left, 0, "a posting of {posting} never settles");
2657 consistent(&ix);
2658 assert_eq!(ix.len(), 160, "settling lost something");
2659 }
2660 }
2661
2662 #[test]
2663 fn a_removed_vector_stops_coming_back() {
2664 let (n, dim, posting) = shrunk(400, 64, Tuning::default().posting);
2665 let store = corpus(dim, n, 4, 11);
2666 let mut ix = build(
2667 &store,
2668 dim,
2669 Tuning {
2670 posting,
2671 ..Tuning::default()
2672 },
2673 );
2674 let q = store.0[7].clone();
2675 assert_eq!(ix.search(&q, 1, &store)[0].id, 7);
2676
2677 assert!(ix.remove(7));
2678 assert!(!ix.remove(7), "removing it twice should say so");
2679 assert!(!ix.contains(7));
2680 assert_eq!(ix.len(), n - 1);
2681 consistent(&ix);
2682 assert!(ix.search(&q, 5, &store).iter().all(|h| h.id != 7));
2683 }
2684
2685 /// How many copies of its members a collection is holding, which is what
2686 /// replication costs and what it has to be paid for in recall.
2687 fn copies(ix: &Partitions) -> f32 {
2688 let held: usize = ix.postings.iter().map(Posting::len).sum();
2689 held as f32 / ix.len() as f32
2690 }
2691
2692 /// The knob does what it says: off means one copy of everything, and on
2693 /// means more than one copy of some things and not of everything.
2694 #[test]
2695 fn spilling_puts_boundary_vectors_in_more_than_one_partition() {
2696 let (n, dim, posting) = shrunk(3000, 32, Tuning::default().posting);
2697 let store = corpus(dim, n, 12, 5);
2698 let base = Tuning {
2699 posting,
2700 ..Tuning::default()
2701 };
2702 let off = Tuning { spill: 1, ..base };
2703 let none = build(&store, dim, off);
2704 consistent(&none);
2705 assert_eq!(copies(&none), 1.0, "spill of one is one copy of everything");
2706
2707 let on = build(&store, dim, base);
2708 consistent(&on);
2709 let rate = copies(&on);
2710 assert!(rate > 1.0, "spilling should make copies, made {rate}");
2711 assert!(
2712 rate < Tuning::default().spill as f32,
2713 "slack should stop short of copying everything into everything, made {rate}"
2714 );
2715 assert_eq!(on.len(), store.0.len(), "a copy is not a member");
2716 }
2717
2718 /// The whole point of it, stated as the thing that is actually true rather
2719 /// than as a recall number.
2720 ///
2721 /// A copy of a member in a second partition means a search that reads that
2722 /// partition finds the member, without widening and without the member's own
2723 /// partition being anywhere near the query. So take a member that got
2724 /// copied, take a different member of the partition it was copied into, and
2725 /// search from that one with a probe of exactly one. The scan reads one
2726 /// posting, and the copy is why the answer is in it.
2727 ///
2728 /// Recall is deliberately not what this asserts. Whether copies pay for
2729 /// themselves end to end is a question about the shape of the data, and on
2730 /// generated vectors the answer is no by a hair, because a tight cluster has
2731 /// no boundary members worth copying and the copies that do get made push
2732 /// the partition count up and the share of the index a fixed probe reads
2733 /// down. `examples/recall.rs` is where that gets answered, on data somebody
2734 /// else made.
2735 #[test]
2736 fn a_copy_is_found_from_the_partition_it_was_copied_into() {
2737 let (n, dim, posting) = shrunk(3000, 32, Tuning::default().posting);
2738 let store = corpus(dim, n, 12, 5);
2739 let t = Tuning {
2740 posting,
2741 slack: 0.25,
2742 ..Tuning::default()
2743 };
2744 let mut ix = build(&store, dim, t);
2745 let (id, copies) = (0..n as u64)
2746 .filter_map(|id| {
2747 let mut places = Vec::new();
2748 ix.every_place(id, &mut places);
2749 (places.len() > 1).then_some((id, places))
2750 })
2751 .next()
2752 .expect("some member near a boundary got copied");
2753
2754 let mut narrow = t;
2755 narrow.probe = 1;
2756 narrow.widen = 1;
2757 ix.retune(narrow);
2758 for place in &copies {
2759 let p = place.partition as usize;
2760 // A different member of the same posting, so the query lands there
2761 // rather than where the copied member belongs.
2762 let neighbour = ix.postings[p]
2763 .ids
2764 .iter()
2765 .copied()
2766 .find(|&other| other != id)
2767 .expect("the partition holds more than the copy");
2768 let got = ix.candidates(&store.0[neighbour as usize], ix.postings[p].len());
2769 assert!(
2770 got.iter().any(|&(seen, _)| seen == id),
2771 "member {id} has a copy in partition {p} and a search of it did not find it"
2772 );
2773 }
2774 }
2775
2776 /// The knob is only worth having if the searches it cuts short were reading
2777 /// partitions that had stopped paying, so the two things to show are that it
2778 /// reads fewer of them and that the answers survive it.
2779 #[test]
2780 #[cfg_attr(
2781 miri,
2782 ignore = "the count is the claim: a saving measured as partitions read per query over a hundred queries, against a recall that has to survive it"
2783 )]
2784 fn patience_reads_fewer_partitions_and_keeps_the_answers() {
2785 let dim = 32;
2786 let store = corpus(dim, 4000, 16, 77);
2787 let wide = Tuning {
2788 probe: 64,
2789 ..Tuning::default()
2790 };
2791 let mut ix = build(&store, dim, wide);
2792 let queries = 100;
2793 let full = recall(&ix, &store, 10, queries);
2794 let cost = |ix: &Partitions| -> f64 {
2795 (0..queries)
2796 .map(|i| {
2797 let q = &store.0[i * 7 % store.0.len()];
2798 ix.search_costed(q, 10, &Any, &store).1.probed
2799 })
2800 .sum::<usize>() as f64
2801 / queries as f64
2802 };
2803 let spent = cost(&ix);
2804
2805 ix.retune(Tuning {
2806 probe: 64,
2807 patience: 2,
2808 ..Tuning::default()
2809 });
2810 let cut = cost(&ix);
2811 assert!(
2812 cut < spent * 0.75,
2813 "patience of two read {cut:.1} partitions a query against {spent:.1}, which is not a saving worth the knob"
2814 );
2815 let after = recall(&ix, &store, 10, queries);
2816 assert!(
2817 after >= full - 0.02,
2818 "recall went from {full} to {after}, which is more than giving up early is allowed to cost"
2819 );
2820 }
2821
2822 /// The rule is written as "once there is enough to answer with", and the
2823 /// case that proves it is the one where there is not. A filter that almost
2824 /// nothing passes is why `widen` exists, and a search that gave up on it
2825 /// after two quiet partitions would return nothing at all.
2826 #[test]
2827 #[cfg_attr(
2828 miri,
2829 ignore = "the count is the claim: ten answers at one member in fifty, spread over enough partitions that the first few cannot hold them"
2830 )]
2831 fn patience_does_not_cut_off_a_filter_that_is_still_short() {
2832 let dim = 32;
2833 let store = corpus(dim, 4000, 16, 91);
2834 let mut ix = build(
2835 &store,
2836 dim,
2837 Tuning {
2838 patience: 1,
2839 ..Tuning::default()
2840 },
2841 );
2842 // One member in fifty, spread over every partition, so the answers are
2843 // certainly not all in the first few.
2844 for id in 0..4000u64 {
2845 ix.retag(id, u64::from(id % 50 == 0));
2846 }
2847 struct Rare;
2848 impl Filter for Rare {
2849 fn allows(&self, tag: u64) -> bool {
2850 tag == 1
2851 }
2852 }
2853 let q = &store.0[3];
2854 let got = ix.search_where(q, 10, &Rare, &store);
2855 assert_eq!(got.len(), 10, "the filtered search came back short");
2856 for hit in &got {
2857 assert!(hit.id.is_multiple_of(50), "{} is not a match", hit.id);
2858 }
2859 }
2860
2861 /// A replicated member is scanned twice by a search that reads both of its
2862 /// partitions, and an answer list with the same id in it twice is a bug the
2863 /// caller sees.
2864 #[test]
2865 fn a_replicated_member_comes_back_once() {
2866 let (n, dim, posting) = shrunk(2000, 32, Tuning::default().posting);
2867 let store = corpus(dim, n, 8, 31);
2868 // Every partition, so that every copy of every member is read and the
2869 // duplicates are certain rather than likely.
2870 let t = Tuning {
2871 posting,
2872 probe: 1 << 20,
2873 ..Tuning::default()
2874 };
2875 let ix = build(&store, dim, t);
2876 for i in 0..50 {
2877 let q = &store.0[i * 37 % store.0.len()];
2878 let got: Vec<u64> = ix.search(q, 20, &store).into_iter().map(|h| h.id).collect();
2879 let mut once = got.clone();
2880 once.sort_unstable();
2881 once.dedup();
2882 assert_eq!(got.len(), once.len(), "a duplicate answer for query {i}");
2883 }
2884 }
2885
2886 /// Every copy has to go, and the arena has to come back. Removing under
2887 /// replication is the path where a leak or a stale placement would show up,
2888 /// and `consistent` is what says it did not.
2889 #[test]
2890 fn removing_a_replicated_member_takes_every_copy() {
2891 let (n, dim, posting) = shrunk(1500, 32, Tuning::default().posting);
2892 let store = corpus(dim, n, 6, 41);
2893 let mut ix = build(
2894 &store,
2895 dim,
2896 Tuning {
2897 posting,
2898 ..Tuning::default()
2899 },
2900 );
2901 let before: usize = ix.postings.iter().map(Posting::len).sum();
2902 let mut gone = 0usize;
2903 for id in (0..n as u64).step_by(3) {
2904 gone += ix.placements_of(id);
2905 assert!(ix.remove(id));
2906 assert!(!ix.contains(id));
2907 }
2908 consistent(&ix);
2909 let after: usize = ix.postings.iter().map(Posting::len).sum();
2910 assert_eq!(before - after, gone, "a copy was left behind");
2911 assert_eq!(ix.len(), n - n.div_ceil(3));
2912 for id in (0..n as u64).step_by(3) {
2913 let q = &store.0[id as usize];
2914 assert!(ix.search(q, 5, &store).iter().all(|h| h.id != id));
2915 }
2916 }
2917
2918 /// A retag has to reach every copy, because a scan meets whichever one it
2919 /// reads first and a filter that sees a stale tag in one partition and a
2920 /// fresh one in another is the worst kind of wrong.
2921 #[test]
2922 fn retagging_a_replicated_member_reaches_every_copy() {
2923 let (n, dim, posting) = shrunk(1200, 32, Tuning::default().posting);
2924 let store = corpus(dim, n, 6, 47);
2925 let tuning = Tuning {
2926 posting,
2927 ..Tuning::default()
2928 };
2929 let mut ix = Partitions::new(dim, Bits::One, 7, tuning);
2930 for (i, v) in store.0.iter().enumerate() {
2931 ix.insert_tagged(i as u64, v, 1);
2932 if i % 64 == 0 {
2933 ix.maintain(&store, 4096);
2934 }
2935 }
2936 ix.maintain(&store, 1 << 20);
2937
2938 // A member put on a boundary rather than one the corpus happened to
2939 // leave there. Halfway between the two centroids nearest each other is
2940 // the same distance from both, which is inside the slack at any width
2941 // and any corpus size, so this is a copy by construction. Scanning for a
2942 // member that spilled on its own worked at twelve hundred vectors in
2943 // thirty two dimensions and found nothing at all at a size Miri can
2944 // afford, which is the usual reward for a test that waits for luck.
2945 let mut pair = (0, 1, f32::INFINITY);
2946 for a in 0..ix.partitions() {
2947 for b in a + 1..ix.partitions() {
2948 let d = sqdist(ix.centroid(a), ix.centroid(b));
2949 if d < pair.2 {
2950 pair = (a, b, d);
2951 }
2952 }
2953 }
2954 let (a, b, _) = pair;
2955 let mid: Vec<f32> = ix
2956 .centroid(a)
2957 .iter()
2958 .zip(ix.centroid(b))
2959 .map(|(x, y)| (x + y) / 2.0)
2960 .collect();
2961 let id = n as u64;
2962 ix.insert_tagged(id, &mid, 1);
2963 assert!(
2964 ix.placements_of(id) > 1,
2965 "a member equidistant from the two nearest centroids was not copied"
2966 );
2967
2968 assert!(ix.retag(id, 9));
2969 let mut copies = Vec::new();
2970 ix.every_place(id, &mut copies);
2971 for place in &copies {
2972 assert_eq!(
2973 ix.postings[place.partition as usize].tags[place.slot as usize], 9,
2974 "a copy kept the old tag"
2975 );
2976 }
2977 consistent(&ix);
2978 }
2979
2980 #[test]
2981 fn inserting_the_same_id_twice_replaces_it() {
2982 let (n, dim, posting) = shrunk(200, 64, Tuning::default().posting);
2983 let store = corpus(dim, n, 2, 13);
2984 let mut ix = build(
2985 &store,
2986 dim,
2987 Tuning {
2988 posting,
2989 ..Tuning::default()
2990 },
2991 );
2992 let before = ix.len();
2993 ix.insert(3, &store.0[3]);
2994 assert_eq!(ix.len(), before);
2995 consistent(&ix);
2996 assert_eq!(ix.search(&store.0[3], 1, &store)[0].id, 3);
2997 }
2998
2999 /// A collection of copies of one vector has no cut in it, and maintenance
3000 /// has to notice that rather than try the same split for ever.
3001 ///
3002 /// This is not a corner case anybody has to go looking for. It is what a
3003 /// collection looks like when a pipeline embeds the same document a thousand
3004 /// times, and getting it wrong is a hang rather than a wrong answer.
3005 #[test]
3006 #[cfg_attr(
3007 miri,
3008 ignore = "the count is the claim: a thousand identical vectors is what makes maintenance try the same split over and over, and the failure it looks for is a hang"
3009 )]
3010 fn a_thousand_copies_of_one_vector_do_not_spin() {
3011 let dim = 32;
3012 let one = corpus(dim, 1, 1, 41).0.pop().expect("one vector");
3013 let store = Store(vec![one; 1000]);
3014 let tuning = Tuning {
3015 posting: 16,
3016 ..Tuning::default()
3017 };
3018 let mut ix = Partitions::new(dim, Bits::One, 7, tuning);
3019 for (i, v) in store.0.iter().enumerate() {
3020 ix.insert(i as u64, v);
3021 ix.maintain(&store, 4096);
3022 }
3023 ix.maintain(&store, 1 << 20);
3024 consistent(&ix);
3025 assert_eq!(ix.len(), 1000);
3026 assert!(!ix.needs_maintenance(), "it still thinks there is work");
3027 // And it still answers, with the exact distance rather than an estimate.
3028 let hits = ix.search(&store.0[0], 5, &store);
3029 assert_eq!(hits.len(), 5);
3030 assert!(hits.iter().all(|h| h.distance < 1e-6));
3031 }
3032
3033 /// G13's actual claim. Recall is measured at the end of a long stream of
3034 /// writes and deletes rather than on a fresh build, because a fresh build is
3035 /// the measurement that hides drift.
3036 #[test]
3037 #[cfg_attr(
3038 miri,
3039 ignore = "the count is the claim: recall at the end of three thousand writes with a tenth of them churned, and a short stream is a fresh build, which is the measurement this one exists to avoid"
3040 )]
3041 fn recall_holds_over_a_write_stream_with_no_rebuild() {
3042 let dim = 96;
3043 let store = corpus(dim, 3000, 15, 17);
3044 let tuning = Tuning {
3045 posting: 64,
3046 ..Tuning::default()
3047 };
3048 let mut ix = Partitions::new(dim, Bits::One, 7, tuning);
3049
3050 // Write everything, and churn a tenth of it as we go, which is what
3051 // moves the centroids around under the members that are already filed.
3052 let mut rng = Rng::new(23);
3053 for (i, v) in store.0.iter().enumerate() {
3054 ix.insert(i as u64, v);
3055 if i > 100 && i % 10 == 0 {
3056 let victim = rng.below(i) as u64;
3057 ix.remove(victim);
3058 ix.insert(victim, &store.0[victim as usize]);
3059 }
3060 ix.maintain(&store, 512);
3061 }
3062 ix.maintain(&store, 1 << 20);
3063 consistent(&ix);
3064 assert_eq!(ix.len(), store.0.len());
3065
3066 let r = recall(&ix, &store, 10, 60);
3067 assert!(r >= 0.95, "recall at 10 after the stream was {r}");
3068 }
3069
3070 /// What the sweep is for, measured as the thing it actually fixes rather
3071 /// than through recall.
3072 ///
3073 /// Drift is members filed under a partition that is no longer their nearest,
3074 /// which is what a split leaves behind in the partitions around it. It shows
3075 /// up in recall eventually, but recall is a blunt instrument here and moves
3076 /// by a percent for reasons that have nothing to do with this, so the
3077 /// straight count is the honest measurement.
3078 #[test]
3079 #[cfg_attr(
3080 miri,
3081 ignore = "the count is the claim: drifted members with the sweep against drifted members without it, and the assertion is the ratio between the two"
3082 )]
3083 fn the_sweep_is_what_keeps_members_under_their_nearest_centroid() {
3084 let dim = 96;
3085 let store = corpus(dim, 2000, 10, 29);
3086 let tuning = Tuning {
3087 posting: 48,
3088 ..Tuning::default()
3089 };
3090 let with = misfiled(&build(&store, dim, tuning), &store);
3091 let without = misfiled(&build(&store, dim, Tuning { sweep: 0, ..tuning }), &store);
3092 assert!(
3093 with * 4 < without,
3094 "sweeping left {with} members drifted and not sweeping left {without}"
3095 );
3096 }
3097
3098 /// How many members are filed under something that is not their nearest
3099 /// centroid.
3100 /// Asked once per member rather than once per posting entry, because a
3101 /// boundary copy sits in a partition that is not the member's nearest on
3102 /// purpose, and counting one as drift would read replication as the very
3103 /// thing the sweep exists to undo. A member has drifted when none of the
3104 /// partitions holding it is its nearest.
3105 fn misfiled(ix: &Partitions, store: &Store) -> usize {
3106 let mut buf = vec![0.0f32; ix.dim()];
3107 let mut wrong = 0;
3108 for id in 0..store.0.len() as u64 {
3109 if !ix.contains(id) {
3110 continue;
3111 }
3112 assert!(store.get(id, &mut buf));
3113 let near = ix.nearest(&ix.quant.rotate(&buf));
3114 if ix.placed_at(id, near).is_none() {
3115 wrong += 1;
3116 }
3117 }
3118 wrong
3119 }
3120
3121 #[test]
3122 fn a_vector_the_log_forgot_is_dropped_rather_than_returned() {
3123 let (n, dim, posting) = shrunk(400, 64, 24);
3124 let store = corpus(dim, n, 4, 31);
3125 let tuning = Tuning {
3126 posting,
3127 ..Tuning::default()
3128 };
3129 let mut ix = build(&store, dim, tuning);
3130 assert!(ix.contains(11));
3131
3132 // The log loses one without telling the index, which is the state a
3133 // crash between two appends leaves behind.
3134 let holey = Holey(store.0.clone(), 11);
3135 assert!(
3136 ix.search(&store.0[11], 5, &holey)
3137 .iter()
3138 .all(|h| h.id != 11)
3139 );
3140
3141 // And maintenance walking over it takes it out for good. Everything but
3142 // the last hundred, so that the postings around the hole get rewritten
3143 // whatever the corpus was.
3144 for id in 0..n as u64 - 100 {
3145 ix.remove(id);
3146 }
3147 ix.maintain(&holey, 1 << 20);
3148 consistent(&ix);
3149 assert!(!ix.contains(11));
3150 }
3151
3152 #[test]
3153 fn rotating_first_is_the_same_as_rotating_inside() {
3154 // The whole index rests on the rotation being linear, so this is the
3155 // property, not an implementation detail.
3156 let dim = 128;
3157 let q = Quantizer::new(dim, Bits::One, 5);
3158 let store = corpus(dim, 2, 1, 37);
3159 let (v, c) = (&store.0[0], &store.0[1]);
3160
3161 let mut a = vec![0u8; q.code_bytes()];
3162 let one = q.encode(v, c, &mut a);
3163 let mut b = vec![0u8; q.code_bytes()];
3164 let two = q.encode_rotated(&q.rotate(v), &q.rotate(c), &mut b);
3165
3166 assert_eq!(a, b, "the two ways round should write the same code");
3167 assert!((one.norm - two.norm).abs() < 1e-4);
3168 assert!((one.scale - two.scale).abs() < 1e-4);
3169 }
3170
3171 #[test]
3172 fn two_means_splits_two_clouds_apart() {
3173 let dim = 8;
3174 let mut xs = Vec::new();
3175 for i in 0..40 {
3176 let far = if i % 2 == 0 { 0.0 } else { 10.0 };
3177 for d in 0..dim {
3178 xs.push(far + (i as f32 + d as f32) * 0.01);
3179 }
3180 }
3181 let (a, b) = two_means(&xs, dim);
3182 let (near, away) = if a[0] < b[0] { (a, b) } else { (b, a) };
3183 assert!(near[0] < 1.0, "{near:?}");
3184 assert!(away[0] > 9.0, "{away:?}");
3185 }
3186}