Skip to main content

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(&centre, 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}