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 thousand copies of the same vector is one point as far as two means
1335        // is concerned, and there is no cut that divides it. Put them back, and
1336        // do not come back until the posting has doubled, so that a collection
1337        // that really is all one vector costs a re-encode of it a logarithmic
1338        // number of times rather than once per insert.
1339        if sides.iter().all(|&s| s) || sides.iter().all(|&s| !s) {
1340            for (i, m) in members.iter().enumerate() {
1341                self.place(p, m.id, m.tag, &xs[i * dim..(i + 1) * dim]);
1342            }
1343            self.postings[p].stuck = members.len() * 2;
1344            return members.len();
1345        }
1346        self.centroids[p * dim..(p + 1) * dim].copy_from_slice(&a);
1347        self.coarse.moved(p, &a, dim);
1348        let q = self.add_partition(&b);
1349        for (i, m) in members.iter().enumerate() {
1350            let to = if sides[i] { p } else { q };
1351            self.place(to, m.id, m.tag, &xs[i * dim..(i + 1) * dim]);
1352        }
1353        members.len() + self.sweep(&[p, q], vectors)
1354    }
1355
1356    /// Hand a partition's members to whoever is nearest now, and drop it.
1357    fn merge(&mut self, p: usize, vectors: &impl Vectors) -> usize {
1358        let (members, xs) = self.take(p, vectors);
1359        let dim = self.dim();
1360        self.drop_partition(p);
1361        // The same approximate lookup an insert uses, and for the same reason.
1362        // An exact one is a distance to every centroid in the collection, once
1363        // per member of the partition being emptied, and a merge is the one
1364        // maintenance job whose whole cost is that lookup. Measured at 1024
1365        // dimensions and 4849 partitions it is most of an ingest.
1366        let mut short = core::mem::take(&mut self.scratch);
1367        for (i, m) in members.iter().enumerate() {
1368            let x = &xs[i * dim..(i + 1) * dim];
1369            let to = self.roughly_nearest(x, &mut short);
1370            self.place(to, m.id, m.tag, x);
1371        }
1372        self.scratch = short;
1373        members.len()
1374    }
1375
1376    /// LIRE: after the centroids move, anything nearby that is now filed under
1377    /// the wrong one gets moved.
1378    ///
1379    /// Only the partitions near the ones that just changed are looked at,
1380    /// because those are the only ones whose members can have a new nearest
1381    /// centroid, and looking at all of them would be the rebuild this index
1382    /// exists to avoid.
1383    ///
1384    /// # Why a member is only measured against what changed
1385    ///
1386    /// Every member is already filed under the centroid it was nearest to, and a
1387    /// split moves one centroid and adds one. Nothing else moved, so for a member
1388    /// of some other partition the nearest of all the centroids that did not
1389    /// change is still the one it is already under, and the only way it can have
1390    /// a new answer is if one of the two new centroids beats that. That is a
1391    /// comparison against two, not a search over all of them.
1392    ///
1393    /// This is not a shortcut, it is what LIRE says, and getting it wrong is
1394    /// expensive in a way that is easy to miss. A sweep after a split walks about
1395    /// four partitions' worth of members, and a split happens every posting's
1396    /// worth of inserts, so a full centroid scan per member works out at several
1397    /// scans of every centroid in the collection per vector inserted. That is the
1398    /// whole ingest cost at any size worth talking about: measured on 128
1399    /// dimensional vectors it was 74 thousand a second at twelve thousand vectors
1400    /// and 13 thousand at two hundred thousand, with maintenance three quarters
1401    /// of it, and `examples/ingest.rs` is the harness that says so.
1402    fn sweep(&mut self, changed: &[usize], vectors: &impl Vectors) -> usize {
1403        let dim = self.dim();
1404        let mut look: Vec<usize> = Vec::new();
1405        for &p in changed {
1406            let centre = self.centroid(p).to_vec();
1407            for q in self.roughly_near_partitions(&centre, self.tuning.sweep) {
1408                if !changed.contains(&q) && !look.contains(&q) {
1409                    look.push(q);
1410                }
1411            }
1412        }
1413        // Copied out because placing a member borrows the index, and safe to
1414        // copy because nothing below here moves a centroid: `place` appends a
1415        // code to a posting and leaves the centroids alone.
1416        let fresh: Vec<(usize, Vec<f32>)> = changed
1417            .iter()
1418            .map(|&p| (p, self.centroid(p).to_vec()))
1419            .collect();
1420        let mut seen = 0;
1421        let mut buf = vec![0.0f32; dim];
1422        for p in look {
1423            let here = self.centroid(p).to_vec();
1424            // Backwards, because taking a member out moves the last one into
1425            // its slot and a backwards walk never steps over the one that moved.
1426            for i in (0..self.postings[p].len()).rev() {
1427                seen += 1;
1428                let id = self.postings[p].ids[i];
1429                let tag = self.postings[p].tags[i];
1430                if !vectors.get(id, &mut buf) {
1431                    self.pull_and_forget(p, i);
1432                    continue;
1433                }
1434                let x = self.quant.rotate(&buf);
1435                let mut best = (p, sqdist(&x, &here));
1436                for (q, centre) in &fresh {
1437                    let d = sqdist(&x, centre);
1438                    if d < best.1 {
1439                        best = (*q, d);
1440                    }
1441                }
1442                if best.0 != p {
1443                    self.pull_and_forget(p, i);
1444                    self.place(best.0, id, tag, &x);
1445                }
1446            }
1447        }
1448        seen
1449    }
1450
1451    /// Empty a partition out, handing back its members and their rotated
1452    /// vectors. Ids the source has forgotten are dropped.
1453    fn take(&mut self, p: usize, vectors: &impl Vectors) -> (Vec<Member>, Vec<f32>) {
1454        let dim = self.dim();
1455        let ids = std::mem::take(&mut self.postings[p].ids);
1456        let tags = std::mem::take(&mut self.postings[p].tags);
1457        self.postings[p].codes.clear();
1458        self.postings[p].meta.clear();
1459        self.note(p);
1460        let mut kept = Vec::with_capacity(ids.len());
1461        let mut xs = Vec::with_capacity(ids.len() * dim);
1462        let mut buf = vec![0.0f32; dim];
1463        for (id, tag) in ids.into_iter().zip(tags) {
1464            // Only this partition's copy. A member replicated into a partition
1465            // that is not the one being emptied keeps the copy it has there.
1466            self.detach(id, p);
1467            if vectors.get(id, &mut buf) {
1468                xs.extend_from_slice(&self.quant.rotate(&buf));
1469                kept.push(Member { id, tag });
1470            }
1471        }
1472        (kept, xs)
1473    }
1474
1475    /// The `n` partitions whose centroids are nearest `x`, nearest first.
1476    ///
1477    /// This measures against every centroid in the collection and it stays that
1478    /// way. Putting it through [`crate::coarse`] was tried and the recall it
1479    /// costs is not worth the time it saves, which the module doc there sets out
1480    /// with the numbers.
1481    /// The `n` partitions near `x`, as far as the coarse layer can tell.
1482    ///
1483    /// This is only for the sweep, and the reason it is allowed there and not on
1484    /// the search path is that it picks which partitions to look in rather than
1485    /// what the answer is. Every member the sweep then looks at is compared
1486    /// exactly against the centroids that just changed, so a neighbour the layer
1487    /// missed costs a few members not moving yet rather than a member moving to
1488    /// the wrong place, and the next split in that neighbourhood picks them up.
1489    /// The distinction matters because an approximate decision inside the sweep
1490    /// is the one thing [`crate::coarse`] says out loud must not happen.
1491    fn roughly_near_partitions(&self, x: &[f32], n: usize) -> Vec<usize> {
1492        if !self.coarse.ready() {
1493            return self.near_partitions(x, n);
1494        }
1495        let mut short = Vec::new();
1496        self.coarse.shortlist(x, self.dim(), &mut short);
1497        let mut by: Vec<(usize, f32)> = short
1498            .iter()
1499            .map(|&p| (p as usize, sqdist(x, self.centroid(p as usize))))
1500            .collect();
1501        let n = n.min(by.len());
1502        by.select_nth_unstable_by(n.saturating_sub(1), |a, b| a.1.total_cmp(&b.1));
1503        by.truncate(n);
1504        by.sort_by(|a, b| a.1.total_cmp(&b.1));
1505        by.into_iter().map(|(p, _)| p).collect()
1506    }
1507
1508    fn near_partitions(&self, x: &[f32], n: usize) -> Vec<usize> {
1509        let mut by: Vec<(usize, f32)> = (0..self.postings.len())
1510            .map(|p| (p, sqdist(x, self.centroid(p))))
1511            .collect();
1512        let n = n.min(by.len());
1513        by.select_nth_unstable_by(n.saturating_sub(1), |a, b| a.1.total_cmp(&b.1));
1514        by.truncate(n);
1515        by.sort_by(|a, b| a.1.total_cmp(&b.1));
1516        by.into_iter().map(|(p, _)| p).collect()
1517    }
1518
1519    /// The partitions in the order a search would probe them, nearest centroid
1520    /// first, all of them.
1521    ///
1522    /// Test only, and it exists for [`miss`](crate::miss), which asks how far
1523    /// down this order a query's true neighbours sit. That is the measurement
1524    /// that says whether the recall gate wants better partitions or a better
1525    /// estimator, and it cannot be taken from outside the crate because the
1526    /// probe order is not something a caller has any business seeing.
1527    #[cfg(test)]
1528    pub(crate) fn probe_order(&self, q: &[f32], into: &mut Vec<usize>) {
1529        let u = self.quant.rotate(q);
1530        *into = self.near_partitions(&u, self.postings.len());
1531    }
1532
1533    /// Which partition holds `id`, if any.
1534    #[cfg(test)]
1535    pub(crate) fn holder(&self, id: u64) -> Option<usize> {
1536        self.any_place(id).map(|s| s.partition as usize)
1537    }
1538
1539    /// The partition `x` belongs to, as far as the coarse layer can tell.
1540    fn roughly_nearest(&self, x: &[f32], short: &mut Vec<u32>) -> usize {
1541        if !self.coarse.ready() {
1542            return self.nearest(x);
1543        }
1544        self.coarse.shortlist(x, self.dim(), short);
1545        short
1546            .iter()
1547            .map(|&p| (p as usize, sqdist(x, self.centroid(p as usize))))
1548            .min_by(|a, b| a.1.total_cmp(&b.1))
1549            .map_or(0, |(p, _)| p)
1550    }
1551
1552    /// The partition `x` belongs to.
1553    fn nearest(&self, x: &[f32]) -> usize {
1554        (0..self.postings.len())
1555            .map(|p| (p, sqdist(x, self.centroid(p))))
1556            .min_by(|a, b| a.1.total_cmp(&b.1))
1557            .map_or(0, |(p, _)| p)
1558    }
1559
1560    fn centroid(&self, p: usize) -> &[f32] {
1561        let dim = self.dim();
1562        &self.centroids[p * dim..(p + 1) * dim]
1563    }
1564
1565    /// A new empty partition around `centroid`, which is already rotated.
1566    fn add_partition(&mut self, centroid: &[f32]) -> usize {
1567        let dim = self.quant.dim();
1568        self.centroids.extend_from_slice(centroid);
1569        self.postings.push(Posting::default());
1570        let p = self.postings.len() - 1;
1571        self.coarse.added(p, centroid, dim);
1572        self.note(p);
1573        self.refresh_coarse();
1574        p
1575    }
1576
1577    /// Rebuild the coarse layer if the partition count has moved far enough
1578    /// since the anchors were last chosen.
1579    fn refresh_coarse(&mut self) {
1580        let n = self.postings.len();
1581        if self.coarse.stale(n) {
1582            let dim = self.quant.dim();
1583            self.coarse.rebuild(&self.centroids, dim, n);
1584        }
1585    }
1586
1587    /// Drop an empty partition, moving the last one into its place.
1588    fn drop_partition(&mut self, p: usize) {
1589        debug_assert_eq!(self.postings[p].len(), 0, "a partition is emptied first");
1590        let dim = self.dim();
1591        let last = self.postings.len() - 1;
1592        self.coarse.dropped(p);
1593        self.postings.swap_remove(p);
1594        for i in 0..dim {
1595            self.centroids[p * dim + i] = self.centroids[last * dim + i];
1596        }
1597        self.centroids.truncate(last * dim);
1598        if p != last {
1599            // The partition that used to be last is at `p` now, so everything
1600            // filed under it has to be told. It is the copy in `last` that
1601            // moves, not the member, so a replicated id keeps its other copies
1602            // pointing where they already point.
1603            for i in 0..self.postings[p].len() {
1604                let id = self.postings[p].ids[i];
1605                if let Some(at) = self.placed_at(id, last) {
1606                    self.places[at as usize].partition = p as u32;
1607                }
1608            }
1609            self.note(p);
1610        }
1611        self.refresh_coarse();
1612    }
1613
1614    /// Append a member to a partition. `x` is rotated.
1615    ///
1616    /// A partition that already holds a copy of `id` keeps the one it has, so
1617    /// that the two maintenance paths that can hand the same member to the same
1618    /// partition twice, a merge into a partition the member was replicated into
1619    /// and a sweep that moves it there, cannot produce a posting with the same
1620    /// id in it twice.
1621    fn place(&mut self, p: usize, id: u64, tag: u64, x: &[f32]) {
1622        let dim = self.dim();
1623        let width = self.quant.code_bytes();
1624        let slot = self.postings[p].len();
1625        if !self.attach(id, p, slot) {
1626            return;
1627        }
1628        self.postings[p].codes.resize((slot + 1) * width, 0);
1629        let centroid = &self.centroids[p * dim..(p + 1) * dim];
1630        let coded = self.quant.encode_rotated(
1631            x,
1632            centroid,
1633            &mut self.postings[p].codes[slot * width..(slot + 1) * width],
1634        );
1635        self.postings[p].ids.push(id);
1636        self.postings[p].tags.push(tag);
1637        self.postings[p].meta.push(coded);
1638        self.note(p);
1639    }
1640
1641    // -- the placement chain -------------------------------------------------
1642    //
1643    // Every site that used to write `self.at` goes through one of these, because
1644    // with replication the question is almost never about an id. It is about one
1645    // copy of an id, the one in a particular partition, and the difference only
1646    // shows up as a corrupt index a long way from where it was caused.
1647
1648    /// Record that `id` has a copy at `(p, slot)`, saying whether it is new.
1649    ///
1650    /// A partition already holding a copy is left alone rather than given a
1651    /// second one. Nothing on the insert path asks for that, but the maintenance
1652    /// paths can: a member replicated into two partitions that are then merged
1653    /// into each other would otherwise arrive twice, and a duplicate inside one
1654    /// posting is the one shape the rest of this cannot cope with, because a
1655    /// delete would take out one copy and leave the other.
1656    fn attach(&mut self, id: u64, p: usize, slot: usize) -> bool {
1657        let head = self.at.get(&id).copied().unwrap_or(END);
1658        let mut walk = head;
1659        while walk != END {
1660            if self.places[walk as usize].partition as usize == p {
1661                return false;
1662            }
1663            walk = self.places[walk as usize].next;
1664        }
1665        let place = Place {
1666            partition: p as u32,
1667            slot: slot as u32,
1668            next: head,
1669        };
1670        let at = if self.free == END {
1671            self.places.push(place);
1672            (self.places.len() - 1) as u32
1673        } else {
1674            let at = self.free;
1675            self.free = self.places[at as usize].next;
1676            self.places[at as usize] = place;
1677            at
1678        };
1679        self.at.insert(id, at);
1680        true
1681    }
1682
1683    /// Forget the copy of `id` in partition `p`, saying whether there was one.
1684    fn detach(&mut self, id: u64, p: usize) -> bool {
1685        let Some(&head) = self.at.get(&id) else {
1686            return false;
1687        };
1688        let mut prev = END;
1689        let mut walk = head;
1690        while walk != END {
1691            let this = self.places[walk as usize];
1692            if this.partition as usize == p {
1693                if prev == END {
1694                    if this.next == END {
1695                        self.at.remove(&id);
1696                    } else {
1697                        self.at.insert(id, this.next);
1698                    }
1699                } else {
1700                    self.places[prev as usize].next = this.next;
1701                }
1702                self.places[walk as usize].next = self.free;
1703                self.free = walk;
1704                return true;
1705            }
1706            prev = walk;
1707            walk = this.next;
1708        }
1709        false
1710    }
1711
1712    /// Forget every copy of `id`, saying whether there were any.
1713    fn detach_all(&mut self, id: u64) -> bool {
1714        let Some(head) = self.at.remove(&id) else {
1715            return false;
1716        };
1717        let mut walk = head;
1718        while walk != END {
1719            let next = self.places[walk as usize].next;
1720            self.places[walk as usize].next = self.free;
1721            self.free = walk;
1722            walk = next;
1723        }
1724        true
1725    }
1726
1727    /// Where the copy of `id` in partition `p` is, if there is one.
1728    fn placed_at(&self, id: u64, p: usize) -> Option<u32> {
1729        let mut walk = self.at.get(&id).copied().unwrap_or(END);
1730        while walk != END {
1731            if self.places[walk as usize].partition as usize == p {
1732                return Some(walk);
1733            }
1734            walk = self.places[walk as usize].next;
1735        }
1736        None
1737    }
1738
1739    /// Say that the copy of `id` in partition `p` is at slot `s` now, which is
1740    /// what a pull leaves behind when it moves the last member into a hole.
1741    fn reslot(&mut self, id: u64, p: usize, s: usize) {
1742        if let Some(at) = self.placed_at(id, p) {
1743            self.places[at as usize].slot = s as u32;
1744        } else {
1745            debug_assert!(
1746                false,
1747                "id {id} is in partition {p} and the map does not say so"
1748            );
1749        }
1750    }
1751
1752    /// How many partitions hold a copy of `id`.
1753    #[cfg(test)]
1754    fn placements_of(&self, id: u64) -> usize {
1755        let mut walk = self.at.get(&id).copied().unwrap_or(END);
1756        let mut n = 0;
1757        while walk != END {
1758            n += 1;
1759            walk = self.places[walk as usize].next;
1760        }
1761        n
1762    }
1763
1764    /// Any one copy of `id`, for the questions that do not care which.
1765    fn any_place(&self, id: u64) -> Option<Place> {
1766        self.at.get(&id).map(|&at| self.places[at as usize])
1767    }
1768
1769    /// Every copy of `id`, collected because the callers that want them all are
1770    /// about to borrow the index mutably.
1771    fn every_place(&self, id: u64, into: &mut Vec<Place>) {
1772        into.clear();
1773        let mut walk = self.at.get(&id).copied().unwrap_or(END);
1774        while walk != END {
1775            let place = self.places[walk as usize];
1776            into.push(place);
1777            walk = place.next;
1778        }
1779    }
1780
1781    /// Take slot `s` out of partition `p`, returning the id that moved into it.
1782    fn pull(&mut self, p: usize, s: usize) -> Option<u64> {
1783        let width = self.quant.code_bytes();
1784        let posting = &mut self.postings[p];
1785        let last = posting.len() - 1;
1786        posting.ids.swap_remove(s);
1787        posting.tags.swap_remove(s);
1788        posting.meta.swap_remove(s);
1789        if s != last {
1790            let (head, tail) = posting.codes.split_at_mut(last * width);
1791            head[s * width..(s + 1) * width].copy_from_slice(&tail[..width]);
1792        }
1793        posting.codes.truncate(last * width);
1794        let moved = (s != last).then(|| posting.ids[s]);
1795        self.note(p);
1796        moved
1797    }
1798
1799    /// The same, keeping the map straight, for the paths that are about to put
1800    /// the member somewhere else.
1801    fn pull_and_forget(&mut self, p: usize, s: usize) {
1802        let id = self.postings[p].ids[s];
1803        self.detach(id, p);
1804        if let Some(moved) = self.pull(p, s) {
1805            self.reslot(moved, p, s);
1806        }
1807    }
1808}
1809
1810/// A member on its way from one partition to another, which is the only time
1811/// its id and its tag travel together without a posting around them.
1812#[derive(Clone, Copy)]
1813struct Member {
1814    id: u64,
1815    tag: u64,
1816}
1817
1818#[derive(Debug, PartialEq, Eq)]
1819enum Job {
1820    Split(usize),
1821    Merge(usize),
1822}
1823
1824/// Two means over a set of vectors laid out end to end.
1825///
1826/// The seeds are the member furthest from the middle and then the member
1827/// furthest from that one, which is deterministic, needs no generator, and
1828/// starts on the axis the cloud is actually longest along. Eight rounds is
1829/// past where this stops moving on anything shaped like an embedding.
1830fn two_means(xs: &[f32], dim: usize) -> (Vec<f32>, Vec<f32>) {
1831    let n = xs.len() / dim;
1832    let mut middle = vec![0.0f32; dim];
1833    for i in 0..n {
1834        for (m, c) in middle.iter_mut().zip(&xs[i * dim..(i + 1) * dim]) {
1835            *m += c;
1836        }
1837    }
1838    for m in &mut middle {
1839        *m /= n as f32;
1840    }
1841    let far = |from: &[f32]| {
1842        (0..n)
1843            .max_by(|&i, &j| {
1844                sqdist(from, &xs[i * dim..(i + 1) * dim])
1845                    .total_cmp(&sqdist(from, &xs[j * dim..(j + 1) * dim]))
1846            })
1847            .unwrap_or(0)
1848    };
1849    let i = far(&middle);
1850    let mut a = xs[i * dim..(i + 1) * dim].to_vec();
1851    let j = far(&a);
1852    let mut b = xs[j * dim..(j + 1) * dim].to_vec();
1853
1854    for _ in 0..8 {
1855        let mut sums = (vec![0.0f32; dim], vec![0.0f32; dim]);
1856        let mut counts = (0usize, 0usize);
1857        for i in 0..n {
1858            let x = &xs[i * dim..(i + 1) * dim];
1859            if sqdist(x, &a) <= sqdist(x, &b) {
1860                for (s, c) in sums.0.iter_mut().zip(x) {
1861                    *s += c;
1862                }
1863                counts.0 += 1;
1864            } else {
1865                for (s, c) in sums.1.iter_mut().zip(x) {
1866                    *s += c;
1867                }
1868                counts.1 += 1;
1869            }
1870        }
1871        // A side that ended up with nothing keeps the seed it had, because a
1872        // mean of no points is not a place and the next round would put every
1873        // member on the other side for ever.
1874        if counts.0 > 0 {
1875            for (m, s) in a.iter_mut().zip(&sums.0) {
1876                *m = s / counts.0 as f32;
1877            }
1878        }
1879        if counts.1 > 0 {
1880            for (m, s) in b.iter_mut().zip(&sums.1) {
1881                *m = s / counts.1 as f32;
1882            }
1883        }
1884    }
1885    (a, b)
1886}
1887
1888/// One candidate, ordered by its estimated distance.
1889///
1890/// The tie break on the id is not decoration. Two members of the same partition
1891/// can get the same estimate out of codes that are 16 bytes wide, and without a
1892/// tie break which of them survives depends on the order the heap happened to
1893/// be in, which makes a search answer depend on the insertion history of the
1894/// collection rather than on the collection.
1895#[derive(PartialEq)]
1896struct Ranked {
1897    at: f32,
1898    id: u64,
1899}
1900
1901impl Eq for Ranked {}
1902
1903impl Ord for Ranked {
1904    fn cmp(&self, other: &Ranked) -> std::cmp::Ordering {
1905        self.at.total_cmp(&other.at).then(self.id.cmp(&other.id))
1906    }
1907}
1908
1909impl PartialOrd for Ranked {
1910    fn partial_cmp(&self, other: &Ranked) -> Option<std::cmp::Ordering> {
1911        Some(self.cmp(other))
1912    }
1913}
1914
1915/// The best `want` candidates seen so far, and nothing else.
1916///
1917/// The scan used to push every member of every partition it read into one
1918/// vector and then select from it, which at probe 64 is 24 thousand entries
1919/// pushed and 24 thousand selected over to keep 160. That is 384 kilobytes of
1920/// writes per search and it was a tenth of the search's time.
1921///
1922/// A bounded heap makes the common case one comparison. Once `want` candidates
1923/// are in, a member is only touched further if it beats the worst of them,
1924/// which after the first partition or two is a small fraction of them, and a
1925/// member that loses never has its id or its tag read at all.
1926struct Bounded {
1927    want: usize,
1928    heap: std::collections::BinaryHeap<Ranked>,
1929}
1930
1931impl Bounded {
1932    fn new(want: usize) -> Bounded {
1933        Bounded {
1934            want,
1935            heap: std::collections::BinaryHeap::with_capacity(want + 1),
1936        }
1937    }
1938
1939    /// Whether there are already `want` answers, which is what says a search
1940    /// that was widening for a filter can stop widening.
1941    fn full(&self) -> bool {
1942        self.heap.len() >= self.want
1943    }
1944
1945    /// Whether `at` could still be one of the answers.
1946    #[inline]
1947    fn wants(&self, at: f32) -> bool {
1948        match self.heap.peek() {
1949            Some(worst) if self.heap.len() >= self.want => at < worst.at,
1950            _ => true,
1951        }
1952    }
1953
1954    fn put(&mut self, id: u64, at: f32) {
1955        if self.heap.len() >= self.want {
1956            self.heap.pop();
1957        }
1958        self.heap.push(Ranked { at, id });
1959    }
1960
1961    /// The answers, nearest first.
1962    fn sorted(self) -> Vec<(u64, f32)> {
1963        self.heap
1964            .into_sorted_vec()
1965            .into_iter()
1966            .map(|r| (r.id, r.at))
1967            .collect()
1968    }
1969}
1970
1971#[cfg(test)]
1972mod tests {
1973    use super::*;
1974    use yo_common::Rng;
1975
1976    /// What `job` used to be: two passes over every partition. The candidate
1977    /// lists are only worth having if they give the same answer, so the slow
1978    /// version stays here as the thing the fast one is checked against.
1979    fn slow_job(ix: &Partitions) -> Option<Job> {
1980        let big = (0..ix.postings.len())
1981            .filter(|&p| ix.postings[p].len() > ix.postings[p].stuck)
1982            .max_by_key(|&p| ix.postings[p].len());
1983        if let Some(big) = big
1984            && ix.postings[big].len() > ix.tuning.posting * 2
1985        {
1986            return Some(Job::Split(big));
1987        }
1988        if ix.postings.len() > 1 {
1989            let small = (0..ix.postings.len()).min_by_key(|&p| ix.postings[p].len())?;
1990            if ix.postings[small].len() * 4 < ix.tuning.posting {
1991                return Some(Job::Merge(small));
1992            }
1993        }
1994        None
1995    }
1996
1997    /// Every partition that qualifies has to be on a list, or maintenance stops
1998    /// happening and the index quietly rots. Deletes are in here on purpose,
1999    /// because a delete is the only thing that pushes a partition down through
2000    /// the merge threshold and it is also what makes a partition disappear and
2001    /// renumber the one that was last.
2002    #[test]
2003    fn the_candidate_lists_answer_what_the_two_passes_answered() {
2004        let store = corpus(16, 3000, 12, 0x105E);
2005        let mut ix = Partitions::new(16, Bits::One, 7, Tuning::default());
2006        let mut rng = Rng::new(0x105F);
2007        let mut live: Vec<u64> = Vec::new();
2008        for id in 0..3000u64 {
2009            ix.insert(id, &store.0[id as usize]);
2010            live.push(id);
2011            if id % 7 == 3 && !live.is_empty() {
2012                let at = rng.below(live.len());
2013                let gone = live.swap_remove(at);
2014                ix.remove(gone);
2015            }
2016            // Once before maintenance runs and once after, because the lists
2017            // are written by both and the state in between is the one a stale
2018            // entry would survive in.
2019            assert_eq!(
2020                ix.job(),
2021                slow_job(&ix),
2022                "after {id} inserts, before maintaining"
2023            );
2024            ix.maintain(&store, 8);
2025            assert_eq!(
2026                ix.job(),
2027                slow_job(&ix),
2028                "after {id} inserts, after maintaining"
2029            );
2030            assert_eq!(ix.needs_maintenance(), slow_job(&ix).is_some());
2031        }
2032        assert!(ix.postings.len() > 5, "the test never split anything");
2033    }
2034
2035    /// The record log, for a test: every vector by id, where the id is where it
2036    /// sits.
2037    struct Store(Vec<Vec<f32>>);
2038
2039    impl Vectors for Store {
2040        fn get(&self, id: u64, into: &mut [f32]) -> bool {
2041            match self.0.get(id as usize) {
2042                Some(v) => {
2043                    into.copy_from_slice(v);
2044                    true
2045                }
2046                None => false,
2047            }
2048        }
2049    }
2050
2051    /// A store with a hole in it, for the case where the log forgot something
2052    /// the index still thinks it has.
2053    struct Holey(Vec<Vec<f32>>, u64);
2054
2055    impl Vectors for Holey {
2056        fn get(&self, id: u64, into: &mut [f32]) -> bool {
2057            if id == self.1 {
2058                return false;
2059            }
2060            match self.0.get(id as usize) {
2061                Some(v) => {
2062                    into.copy_from_slice(v);
2063                    true
2064                }
2065                None => false,
2066            }
2067        }
2068    }
2069
2070    /// Vectors with the two things real embeddings have and uniform noise does
2071    /// not: a few coordinates carrying most of the energy, so that two vectors
2072    /// are genuinely near each other rather than all being equally far apart,
2073    /// and clusters, so that the partitioning has something to get right.
2074    ///
2075    /// A corpus without the first of those is not a hard test, it is an
2076    /// impossible one. The nearest ten of three hundred uniform points are
2077    /// arbitrary, no quantiser can pick them out, and the recall it measures
2078    /// says nothing about the index.
2079    fn corpus(dim: usize, n: usize, clusters: usize, seed: u64) -> Store {
2080        let mut rng = Rng::new(seed);
2081        let centres: Vec<Vec<f32>> = (0..clusters).map(|_| draw(dim, &mut rng)).collect();
2082        Store(
2083            (0..n)
2084                .map(|i| {
2085                    let off = draw(dim, &mut rng);
2086                    let mut v: Vec<f32> = centres[i % clusters]
2087                        .iter()
2088                        .zip(&off)
2089                        .map(|(c, o)| c + o * 0.7)
2090                        .collect();
2091                    unit(&mut v);
2092                    v
2093                })
2094                .collect(),
2095        )
2096    }
2097
2098    /// One vector of the shape above, unit length.
2099    fn draw(dim: usize, rng: &mut Rng) -> Vec<f32> {
2100        let mut v: Vec<f32> = (0..dim)
2101            .map(|i| {
2102                let u = (rng.next_u64() >> 40) as f32 / (1u32 << 24) as f32;
2103                let heavy = if i < dim / 16 { 6.0 } else { 1.0 };
2104                (u * 2.0 - 1.0) * heavy
2105            })
2106            .collect();
2107        unit(&mut v);
2108        v
2109    }
2110
2111    fn unit(v: &mut [f32]) {
2112        let len = v.iter().map(|c| c * c).sum::<f32>().sqrt();
2113        for c in v {
2114            *c /= len;
2115        }
2116    }
2117
2118    fn truth(store: &Store, q: &[f32], k: usize) -> Vec<u64> {
2119        let mut all: Vec<(u64, f32)> = store
2120            .0
2121            .iter()
2122            .enumerate()
2123            .map(|(i, v)| (i as u64, sqdist(q, v)))
2124            .collect();
2125        all.sort_by(|a, b| a.1.total_cmp(&b.1));
2126        all.truncate(k);
2127        all.into_iter().map(|(i, _)| i).collect()
2128    }
2129
2130    /// Build an index over a whole store, running maintenance as it goes the way
2131    /// a maintenance slice would.
2132    fn build(store: &Store, dim: usize, tuning: Tuning) -> Partitions {
2133        let mut ix = Partitions::new(dim, Bits::One, 7, tuning);
2134        for (i, v) in store.0.iter().enumerate() {
2135            ix.insert(i as u64, v);
2136            if i % 64 == 0 {
2137                ix.maintain(store, 4096);
2138            }
2139        }
2140        ix.maintain(store, 1 << 20);
2141        ix
2142    }
2143
2144    /// How often the true `k` nearest come back.
2145    fn recall(ix: &Partitions, store: &Store, k: usize, queries: usize) -> f32 {
2146        let mut hits = 0usize;
2147        for i in 0..queries {
2148            // A query near a real vector rather than anywhere at all, because
2149            // that is what a search looks like.
2150            let q = &store.0[i * 7 % store.0.len()];
2151            let want = truth(store, q, k);
2152            let got: Vec<u64> = ix.search(q, k, store).into_iter().map(|h| h.id).collect();
2153            hits += want.iter().filter(|id| got.contains(id)).count();
2154        }
2155        hits as f32 / (queries * k) as f32
2156    }
2157
2158    /// Everything the index believes about itself, checked.
2159    fn consistent(ix: &Partitions) {
2160        assert_eq!(ix.centroids.len(), ix.postings.len() * ix.dim());
2161        let width = ix.quant.code_bytes();
2162        let mut seen = 0usize;
2163        for (p, posting) in ix.postings.iter().enumerate() {
2164            assert_eq!(posting.codes.len(), posting.len() * width, "partition {p}");
2165            assert_eq!(posting.meta.len(), posting.len(), "partition {p}");
2166            assert_eq!(posting.tags.len(), posting.len(), "partition {p}");
2167            let mut here = HashSet::new();
2168            for (s, id) in posting.ids.iter().enumerate() {
2169                assert!(here.insert(*id), "id {id} is twice in partition {p}");
2170                let at = ix
2171                    .placed_at(*id, p)
2172                    .expect("every member is in the map, under the partition holding it");
2173                assert_eq!(ix.places[at as usize].slot as usize, s, "id {id}");
2174                seen += 1;
2175            }
2176        }
2177        // Every placement points at a member, as many placements as there are
2178        // members, and the free list accounts for the rest of the arena. A
2179        // replicated id makes the first of those the interesting one: a chain
2180        // that kept an entry for a copy that was pulled would still look right
2181        // from the posting's side, and the count is what catches it.
2182        let mut held = 0usize;
2183        for (&id, &head) in &ix.at {
2184            let mut walk = head;
2185            let mut mine = HashSet::new();
2186            while walk != END {
2187                let place = ix.places[walk as usize];
2188                let p = place.partition as usize;
2189                assert!(mine.insert(p), "id {id} is filed twice under partition {p}");
2190                assert!(
2191                    p < ix.postings.len(),
2192                    "id {id} is filed under partition {p}"
2193                );
2194                assert_eq!(
2195                    ix.postings[p].ids[place.slot as usize], id,
2196                    "id {id} is filed at a slot holding something else"
2197                );
2198                held += 1;
2199                walk = place.next;
2200            }
2201        }
2202        assert_eq!(seen, held, "the map and the postings disagree on the count");
2203        let mut spare = 0usize;
2204        let mut walk = ix.free;
2205        while walk != END {
2206            spare += 1;
2207            assert!(spare <= ix.places.len(), "the free list has a cycle in it");
2208            walk = ix.places[walk as usize].next;
2209        }
2210        assert_eq!(held + spare, ix.places.len(), "the arena has leaked");
2211    }
2212
2213    /// Build an index where every vector carries a tag, so the filter has
2214    /// something to meet.
2215    fn build_tagged(
2216        store: &Store,
2217        dim: usize,
2218        tuning: Tuning,
2219        tag: impl Fn(u64) -> u64,
2220    ) -> Partitions {
2221        let mut ix = Partitions::new(dim, Bits::One, 7, tuning);
2222        for (i, v) in store.0.iter().enumerate() {
2223            ix.insert_tagged(i as u64, v, tag(i as u64));
2224            if i % 64 == 0 {
2225                ix.maintain(store, 4096);
2226            }
2227        }
2228        ix.maintain(store, 1 << 20);
2229        ix
2230    }
2231
2232    /// The whole point of pushing a filter into the scan, measured against the
2233    /// thing it replaces.
2234    ///
2235    /// One document in fifty carries the tag. Filtering inside the scan finds
2236    /// the true ten of those. Taking the best forty by vector and then throwing
2237    /// away the ones that do not match, which is what a search that cannot push
2238    /// a filter down has to do, finds almost none of them, and the ones it
2239    /// misses were nearer than the ones it kept.
2240    #[test]
2241    fn a_filter_in_the_scan_finds_what_a_filter_after_it_cannot() {
2242        let dim = 96;
2243        let store = corpus(dim, 3000, 12, 47);
2244        let tuning = Tuning {
2245            posting: 64,
2246            ..Tuning::default()
2247        };
2248        let wanted = |id: u64| id.is_multiple_of(50);
2249        let ix = build_tagged(&store, dim, tuning, |id| u64::from(wanted(id)));
2250
2251        let (mut pushed, mut after) = (0usize, 0usize);
2252        let k = 10;
2253        for i in 0..40 {
2254            let q = &store.0[i * 71 % store.0.len()];
2255
2256            // What the answer is: brute force over the members that match.
2257            let mut all: Vec<(u64, f32)> = store
2258                .0
2259                .iter()
2260                .enumerate()
2261                .filter(|(id, _)| wanted(*id as u64))
2262                .map(|(id, v)| (id as u64, sqdist(q, v)))
2263                .collect();
2264            all.sort_by(|a, b| a.1.total_cmp(&b.1));
2265            let want: Vec<u64> = all[..k].iter().map(|(id, _)| *id).collect();
2266
2267            let got: Vec<u64> = ix
2268                .search_where(q, k, &|tag: u64| tag == 1, &store)
2269                .into_iter()
2270                .map(|h| h.id)
2271                .collect();
2272            pushed += want.iter().filter(|id| got.contains(id)).count();
2273
2274            let late: Vec<u64> = ix
2275                .search(q, k * tuning.rerank, &store)
2276                .into_iter()
2277                .map(|h| h.id)
2278                .filter(|id| wanted(*id))
2279                .take(k)
2280                .collect();
2281            after += want.iter().filter(|id| late.contains(id)).count();
2282        }
2283        let (pushed, after) = (pushed as f32 / 400.0, after as f32 / 400.0);
2284        assert!(pushed >= 0.95, "pushing the filter down gave {pushed}");
2285        assert!(
2286            after < pushed / 2.0,
2287            "filtering afterwards gave {after} against {pushed}, which is not the point being made"
2288        );
2289    }
2290
2291    /// The second test decides, and the scan keeps widening until it has `k` of
2292    /// what the second test wants rather than `k` of what the tag wants.
2293    ///
2294    /// This is what a filter whose tag is only a summary looks like: the tag
2295    /// lets a superset through, so an answer that passes it and fails the exact
2296    /// test must not have cost an answer that passes both.
2297    #[test]
2298    fn the_exact_test_decides_and_the_scan_widens_for_it() {
2299        struct Summary;
2300
2301        impl Filter for Summary {
2302            fn allows(&self, tag: u64) -> bool {
2303                // One in ten, which is what a bit that several values landed on
2304                // looks like from inside the scan.
2305                tag == 1
2306            }
2307
2308            fn exact(&self, id: u64) -> bool {
2309                // One in fifty, and a subset of what the tag allowed, which is
2310                // the direction a summary is allowed to be wrong in.
2311                id.is_multiple_of(50)
2312            }
2313        }
2314
2315        let dim = 64;
2316        let store = corpus(dim, 3000, 9, 71);
2317        let tuning = Tuning {
2318            posting: 64,
2319            ..Tuning::default()
2320        };
2321        let ix = build_tagged(&store, dim, tuning, |id| u64::from(id.is_multiple_of(10)));
2322
2323        let k = 10;
2324        let mut found = 0usize;
2325        for i in 0..20 {
2326            let q = &store.0[i * 131 % store.0.len()];
2327            let mut all: Vec<(u64, f32)> = store
2328                .0
2329                .iter()
2330                .enumerate()
2331                .map(|(id, v)| (id as u64, sqdist(q, v)))
2332                .filter(|(id, _)| id.is_multiple_of(50))
2333                .collect();
2334            all.sort_by(|a, b| a.1.total_cmp(&b.1));
2335            let want: Vec<u64> = all[..k].iter().map(|(id, _)| *id).collect();
2336
2337            let got: Vec<u64> = ix
2338                .search_where(q, k, &Summary, &store)
2339                .into_iter()
2340                .map(|h| h.id)
2341                .collect();
2342            assert!(
2343                got.iter().all(|id| id.is_multiple_of(50)),
2344                "the exact test did not decide: {got:?}"
2345            );
2346            found += want.iter().filter(|id| got.contains(id)).count();
2347        }
2348        let recall = found as f32 / (20.0 * k as f32);
2349        assert!(recall >= 0.9, "two stage filtering gave {recall}");
2350    }
2351
2352    #[test]
2353    fn a_filter_that_matches_nothing_answers_nothing() {
2354        let dim = 64;
2355        let store = corpus(dim, 500, 4, 53);
2356        let ix = build_tagged(&store, dim, Tuning::default(), |_| 1);
2357        assert!(
2358            ix.search_where(&store.0[0], 10, &|tag: u64| tag == 2, &store)
2359                .is_empty()
2360        );
2361        // And the same filter matching everything is the unfiltered answer.
2362        let all = ix.search_where(&store.0[0], 10, &Any, &store);
2363        assert_eq!(all, ix.search(&store.0[0], 10, &store));
2364    }
2365
2366    #[test]
2367    fn a_tag_survives_a_split_and_a_merge() {
2368        let dim = 64;
2369        let store = corpus(dim, 800, 6, 59);
2370        let tuning = Tuning {
2371            posting: 24,
2372            ..Tuning::default()
2373        };
2374        let mut ix = build_tagged(&store, dim, tuning, |id| id * 7 + 1);
2375        assert!(ix.partitions() > 4, "it never split");
2376        for id in 0..800u64 {
2377            assert_eq!(ix.tag(id), Some(id * 7 + 1), "id {id} after the splits");
2378        }
2379
2380        // Now shrink it until partitions merge, and the survivors keep theirs.
2381        for id in 0..760u64 {
2382            ix.remove(id);
2383        }
2384        ix.maintain(&store, 1 << 20);
2385        consistent(&ix);
2386        for id in 760..800u64 {
2387            assert_eq!(ix.tag(id), Some(id * 7 + 1), "id {id} after the merges");
2388        }
2389        assert_eq!(ix.tag(0), None);
2390    }
2391
2392    /// A selective filter means the answers are not in the nearest partitions,
2393    /// and a search that will not look further returns fewer than it should.
2394    #[test]
2395    fn a_selective_filter_makes_the_search_look_further() {
2396        let dim = 64;
2397        let store = corpus(dim, 2000, 10, 61);
2398        let tuning = Tuning {
2399            posting: 32,
2400            ..Tuning::default()
2401        };
2402        let tag = |id: u64| u64::from(id.is_multiple_of(100));
2403        let ix = build_tagged(&store, dim, tuning, tag);
2404        let narrow = build_tagged(&store, dim, Tuning { widen: 1, ..tuning }, tag);
2405
2406        let mut wide_found = 0usize;
2407        let mut narrow_found = 0usize;
2408        for i in 0..20 {
2409            let q = &store.0[i * 91 % store.0.len()];
2410            wide_found += ix.search_where(q, 10, &|t: u64| t == 1, &store).len();
2411            narrow_found += narrow.search_where(q, 10, &|t: u64| t == 1, &store).len();
2412        }
2413        assert_eq!(
2414            wide_found, 200,
2415            "one in a hundred of two thousand is twenty"
2416        );
2417        assert!(
2418            narrow_found < wide_found,
2419            "not widening found {narrow_found} of {wide_found}"
2420        );
2421    }
2422
2423    #[test]
2424    fn a_signature_never_rejects_something_it_should_have_matched() {
2425        let english = Signature::of(&[("lang", b"en")]);
2426        let doc = Signature::of(&[("lang", b"en"), ("topic", b"finance"), ("year", b"2026")]);
2427        assert!(doc.covers(english));
2428        assert!(english.allows(doc.bits()));
2429        assert_eq!(Signature::from_bits(doc.bits()), doc);
2430
2431        // And over a lot of values, nothing that matches is ever turned away.
2432        for i in 0..500u32 {
2433            let value = i.to_string();
2434            let one = Signature::of(&[("id", value.as_bytes())]);
2435            let with = Signature::of(&[("id", value.as_bytes()), ("kind", b"page")]);
2436            assert!(with.covers(one), "value {value}");
2437        }
2438    }
2439
2440    #[test]
2441    fn an_empty_index_answers_nothing() {
2442        let ix = Partitions::new(32, Bits::One, 1, Tuning::default());
2443        let store = Store(Vec::new());
2444        assert!(ix.is_empty());
2445        assert_eq!(ix.partitions(), 0);
2446        assert!(ix.search(&[0.0; 32], 10, &store).is_empty());
2447        assert!(!ix.needs_maintenance());
2448    }
2449
2450    #[test]
2451    fn the_first_vector_is_the_first_partition() {
2452        let store = corpus(32, 1, 1, 3);
2453        let mut ix = Partitions::new(32, Bits::One, 1, Tuning::default());
2454        ix.insert(0, &store.0[0]);
2455        assert_eq!(ix.partitions(), 1);
2456        assert_eq!(ix.len(), 1);
2457        let hits = ix.search(&store.0[0], 5, &store);
2458        assert_eq!(hits.len(), 1);
2459        assert_eq!(hits[0].id, 0);
2460        assert!(hits[0].distance < 1e-6, "{}", hits[0].distance);
2461        consistent(&ix);
2462    }
2463
2464    #[test]
2465    fn a_search_finds_what_brute_force_finds() {
2466        let dim = 128;
2467        let store = corpus(dim, 2000, 12, 5);
2468        let ix = build(&store, dim, Tuning::default());
2469        assert!(ix.partitions() > 1, "it never split");
2470        consistent(&ix);
2471        let r = recall(&ix, &store, 10, 50);
2472        assert!(r >= 0.95, "recall at 10 was {r}");
2473    }
2474
2475    #[test]
2476    fn a_posting_that_grows_too_big_splits() {
2477        let dim = 64;
2478        let tuning = Tuning {
2479            posting: 32,
2480            ..Tuning::default()
2481        };
2482        let store = corpus(dim, 600, 6, 9);
2483        let ix = build(&store, dim, tuning);
2484        assert!(
2485            ix.partitions() >= 600 / (32 * 2),
2486            "600 vectors in {} partitions",
2487            ix.partitions()
2488        );
2489        for posting in &ix.postings {
2490            assert!(
2491                posting.len() <= 32 * 2,
2492                "a posting is {} long",
2493                posting.len()
2494            );
2495        }
2496        consistent(&ix);
2497    }
2498
2499    #[test]
2500    fn a_posting_that_shrinks_merges() {
2501        let dim = 64;
2502        let tuning = Tuning {
2503            posting: 32,
2504            ..Tuning::default()
2505        };
2506        let store = corpus(dim, 600, 6, 9);
2507        let mut ix = build(&store, dim, tuning);
2508        let grown = ix.partitions();
2509        assert!(grown > 4);
2510
2511        // Take away almost everything and let maintenance settle.
2512        for id in 0..570u64 {
2513            assert!(ix.remove(id));
2514        }
2515        ix.maintain(&store, 1 << 20);
2516        consistent(&ix);
2517        assert_eq!(ix.len(), 30);
2518        assert!(
2519            ix.partitions() < grown,
2520            "{} partitions for 30 vectors, was {grown}",
2521            ix.partitions()
2522        );
2523        // And it still answers.
2524        let hits = ix.search(&store.0[599], 1, &store);
2525        assert_eq!(hits[0].id, 599);
2526    }
2527
2528    #[test]
2529    fn a_removed_vector_stops_coming_back() {
2530        let dim = 64;
2531        let store = corpus(dim, 400, 4, 11);
2532        let mut ix = build(&store, dim, Tuning::default());
2533        let q = store.0[7].clone();
2534        assert_eq!(ix.search(&q, 1, &store)[0].id, 7);
2535
2536        assert!(ix.remove(7));
2537        assert!(!ix.remove(7), "removing it twice should say so");
2538        assert!(!ix.contains(7));
2539        assert_eq!(ix.len(), 399);
2540        consistent(&ix);
2541        assert!(ix.search(&q, 5, &store).iter().all(|h| h.id != 7));
2542    }
2543
2544    /// How many copies of its members a collection is holding, which is what
2545    /// replication costs and what it has to be paid for in recall.
2546    fn copies(ix: &Partitions) -> f32 {
2547        let held: usize = ix.postings.iter().map(Posting::len).sum();
2548        held as f32 / ix.len() as f32
2549    }
2550
2551    /// The knob does what it says: off means one copy of everything, and on
2552    /// means more than one copy of some things and not of everything.
2553    #[test]
2554    fn spilling_puts_boundary_vectors_in_more_than_one_partition() {
2555        let dim = 32;
2556        let store = corpus(dim, 3000, 12, 5);
2557        let off = Tuning {
2558            spill: 1,
2559            ..Tuning::default()
2560        };
2561        let none = build(&store, dim, off);
2562        consistent(&none);
2563        assert_eq!(copies(&none), 1.0, "spill of one is one copy of everything");
2564
2565        let on = build(&store, dim, Tuning::default());
2566        consistent(&on);
2567        let rate = copies(&on);
2568        assert!(rate > 1.0, "spilling should make copies, made {rate}");
2569        assert!(
2570            rate < Tuning::default().spill as f32,
2571            "slack should stop short of copying everything into everything, made {rate}"
2572        );
2573        assert_eq!(on.len(), store.0.len(), "a copy is not a member");
2574    }
2575
2576    /// The whole point of it, stated as the thing that is actually true rather
2577    /// than as a recall number.
2578    ///
2579    /// A copy of a member in a second partition means a search that reads that
2580    /// partition finds the member, without widening and without the member's own
2581    /// partition being anywhere near the query. So take a member that got
2582    /// copied, take a different member of the partition it was copied into, and
2583    /// search from that one with a probe of exactly one. The scan reads one
2584    /// posting, and the copy is why the answer is in it.
2585    ///
2586    /// Recall is deliberately not what this asserts. Whether copies pay for
2587    /// themselves end to end is a question about the shape of the data, and on
2588    /// generated vectors the answer is no by a hair, because a tight cluster has
2589    /// no boundary members worth copying and the copies that do get made push
2590    /// the partition count up and the share of the index a fixed probe reads
2591    /// down. `examples/recall.rs` is where that gets answered, on data somebody
2592    /// else made.
2593    #[test]
2594    fn a_copy_is_found_from_the_partition_it_was_copied_into() {
2595        let dim = 32;
2596        let store = corpus(dim, 3000, 12, 5);
2597        let t = Tuning {
2598            slack: 0.25,
2599            ..Tuning::default()
2600        };
2601        let mut ix = build(&store, dim, t);
2602        let (id, copies) = (0..3000u64)
2603            .filter_map(|id| {
2604                let mut places = Vec::new();
2605                ix.every_place(id, &mut places);
2606                (places.len() > 1).then_some((id, places))
2607            })
2608            .next()
2609            .expect("some member near a boundary got copied");
2610
2611        let mut narrow = t;
2612        narrow.probe = 1;
2613        narrow.widen = 1;
2614        ix.retune(narrow);
2615        for place in &copies {
2616            let p = place.partition as usize;
2617            // A different member of the same posting, so the query lands there
2618            // rather than where the copied member belongs.
2619            let neighbour = ix.postings[p]
2620                .ids
2621                .iter()
2622                .copied()
2623                .find(|&other| other != id)
2624                .expect("the partition holds more than the copy");
2625            let got = ix.candidates(&store.0[neighbour as usize], ix.postings[p].len());
2626            assert!(
2627                got.iter().any(|&(seen, _)| seen == id),
2628                "member {id} has a copy in partition {p} and a search of it did not find it"
2629            );
2630        }
2631    }
2632
2633    /// The knob is only worth having if the searches it cuts short were reading
2634    /// partitions that had stopped paying, so the two things to show are that it
2635    /// reads fewer of them and that the answers survive it.
2636    #[test]
2637    fn patience_reads_fewer_partitions_and_keeps_the_answers() {
2638        let dim = 32;
2639        let store = corpus(dim, 4000, 16, 77);
2640        let wide = Tuning {
2641            probe: 64,
2642            ..Tuning::default()
2643        };
2644        let mut ix = build(&store, dim, wide);
2645        let queries = 100;
2646        let full = recall(&ix, &store, 10, queries);
2647        let cost = |ix: &Partitions| -> f64 {
2648            (0..queries)
2649                .map(|i| {
2650                    let q = &store.0[i * 7 % store.0.len()];
2651                    ix.search_costed(q, 10, &Any, &store).1.probed
2652                })
2653                .sum::<usize>() as f64
2654                / queries as f64
2655        };
2656        let spent = cost(&ix);
2657
2658        ix.retune(Tuning {
2659            probe: 64,
2660            patience: 2,
2661            ..Tuning::default()
2662        });
2663        let cut = cost(&ix);
2664        assert!(
2665            cut < spent * 0.75,
2666            "patience of two read {cut:.1} partitions a query against {spent:.1}, which is not a saving worth the knob"
2667        );
2668        let after = recall(&ix, &store, 10, queries);
2669        assert!(
2670            after >= full - 0.02,
2671            "recall went from {full} to {after}, which is more than giving up early is allowed to cost"
2672        );
2673    }
2674
2675    /// The rule is written as "once there is enough to answer with", and the
2676    /// case that proves it is the one where there is not. A filter that almost
2677    /// nothing passes is why `widen` exists, and a search that gave up on it
2678    /// after two quiet partitions would return nothing at all.
2679    #[test]
2680    fn patience_does_not_cut_off_a_filter_that_is_still_short() {
2681        let dim = 32;
2682        let store = corpus(dim, 4000, 16, 91);
2683        let mut ix = build(
2684            &store,
2685            dim,
2686            Tuning {
2687                patience: 1,
2688                ..Tuning::default()
2689            },
2690        );
2691        // One member in fifty, spread over every partition, so the answers are
2692        // certainly not all in the first few.
2693        for id in 0..4000u64 {
2694            ix.retag(id, u64::from(id % 50 == 0));
2695        }
2696        struct Rare;
2697        impl Filter for Rare {
2698            fn allows(&self, tag: u64) -> bool {
2699                tag == 1
2700            }
2701        }
2702        let q = &store.0[3];
2703        let got = ix.search_where(q, 10, &Rare, &store);
2704        assert_eq!(got.len(), 10, "the filtered search came back short");
2705        for hit in &got {
2706            assert!(hit.id.is_multiple_of(50), "{} is not a match", hit.id);
2707        }
2708    }
2709
2710    /// A replicated member is scanned twice by a search that reads both of its
2711    /// partitions, and an answer list with the same id in it twice is a bug the
2712    /// caller sees.
2713    #[test]
2714    fn a_replicated_member_comes_back_once() {
2715        let dim = 32;
2716        let store = corpus(dim, 2000, 8, 31);
2717        // Every partition, so that every copy of every member is read and the
2718        // duplicates are certain rather than likely.
2719        let t = Tuning {
2720            probe: 1 << 20,
2721            ..Tuning::default()
2722        };
2723        let ix = build(&store, dim, t);
2724        for i in 0..50 {
2725            let q = &store.0[i * 37 % store.0.len()];
2726            let got: Vec<u64> = ix.search(q, 20, &store).into_iter().map(|h| h.id).collect();
2727            let mut once = got.clone();
2728            once.sort_unstable();
2729            once.dedup();
2730            assert_eq!(got.len(), once.len(), "a duplicate answer for query {i}");
2731        }
2732    }
2733
2734    /// Every copy has to go, and the arena has to come back. Removing under
2735    /// replication is the path where a leak or a stale placement would show up,
2736    /// and `consistent` is what says it did not.
2737    #[test]
2738    fn removing_a_replicated_member_takes_every_copy() {
2739        let dim = 32;
2740        let store = corpus(dim, 1500, 6, 41);
2741        let mut ix = build(&store, dim, Tuning::default());
2742        let before: usize = ix.postings.iter().map(Posting::len).sum();
2743        let mut gone = 0usize;
2744        for id in (0..1500u64).step_by(3) {
2745            gone += ix.placements_of(id);
2746            assert!(ix.remove(id));
2747            assert!(!ix.contains(id));
2748        }
2749        consistent(&ix);
2750        let after: usize = ix.postings.iter().map(Posting::len).sum();
2751        assert_eq!(before - after, gone, "a copy was left behind");
2752        assert_eq!(ix.len(), 1000);
2753        for id in (0..1500u64).step_by(3) {
2754            let q = &store.0[id as usize];
2755            assert!(ix.search(q, 5, &store).iter().all(|h| h.id != id));
2756        }
2757    }
2758
2759    /// A retag has to reach every copy, because a scan meets whichever one it
2760    /// reads first and a filter that sees a stale tag in one partition and a
2761    /// fresh one in another is the worst kind of wrong.
2762    #[test]
2763    fn retagging_a_replicated_member_reaches_every_copy() {
2764        let dim = 32;
2765        let store = corpus(dim, 1200, 6, 47);
2766        let mut ix = Partitions::new(dim, Bits::One, 7, Tuning::default());
2767        for (i, v) in store.0.iter().enumerate() {
2768            ix.insert_tagged(i as u64, v, 1);
2769            if i % 64 == 0 {
2770                ix.maintain(&store, 4096);
2771            }
2772        }
2773        ix.maintain(&store, 1 << 20);
2774        let spread = (0..1200u64).find(|&id| ix.placements_of(id) > 1);
2775        let id = spread.expect("some member is in more than one partition");
2776        assert!(ix.retag(id, 9));
2777        let mut copies = Vec::new();
2778        ix.every_place(id, &mut copies);
2779        for place in &copies {
2780            assert_eq!(
2781                ix.postings[place.partition as usize].tags[place.slot as usize], 9,
2782                "a copy kept the old tag"
2783            );
2784        }
2785        consistent(&ix);
2786    }
2787
2788    #[test]
2789    fn inserting_the_same_id_twice_replaces_it() {
2790        let dim = 64;
2791        let store = corpus(dim, 200, 2, 13);
2792        let mut ix = build(&store, dim, Tuning::default());
2793        let before = ix.len();
2794        ix.insert(3, &store.0[3]);
2795        assert_eq!(ix.len(), before);
2796        consistent(&ix);
2797        assert_eq!(ix.search(&store.0[3], 1, &store)[0].id, 3);
2798    }
2799
2800    /// A collection of copies of one vector has no cut in it, and maintenance
2801    /// has to notice that rather than try the same split for ever.
2802    ///
2803    /// This is not a corner case anybody has to go looking for. It is what a
2804    /// collection looks like when a pipeline embeds the same document a thousand
2805    /// times, and getting it wrong is a hang rather than a wrong answer.
2806    #[test]
2807    fn a_thousand_copies_of_one_vector_do_not_spin() {
2808        let dim = 32;
2809        let one = corpus(dim, 1, 1, 41).0.pop().expect("one vector");
2810        let store = Store(vec![one; 1000]);
2811        let tuning = Tuning {
2812            posting: 16,
2813            ..Tuning::default()
2814        };
2815        let mut ix = Partitions::new(dim, Bits::One, 7, tuning);
2816        for (i, v) in store.0.iter().enumerate() {
2817            ix.insert(i as u64, v);
2818            ix.maintain(&store, 4096);
2819        }
2820        ix.maintain(&store, 1 << 20);
2821        consistent(&ix);
2822        assert_eq!(ix.len(), 1000);
2823        assert!(!ix.needs_maintenance(), "it still thinks there is work");
2824        // And it still answers, with the exact distance rather than an estimate.
2825        let hits = ix.search(&store.0[0], 5, &store);
2826        assert_eq!(hits.len(), 5);
2827        assert!(hits.iter().all(|h| h.distance < 1e-6));
2828    }
2829
2830    /// G13's actual claim. Recall is measured at the end of a long stream of
2831    /// writes and deletes rather than on a fresh build, because a fresh build is
2832    /// the measurement that hides drift.
2833    #[test]
2834    fn recall_holds_over_a_write_stream_with_no_rebuild() {
2835        let dim = 96;
2836        let store = corpus(dim, 3000, 15, 17);
2837        let tuning = Tuning {
2838            posting: 64,
2839            ..Tuning::default()
2840        };
2841        let mut ix = Partitions::new(dim, Bits::One, 7, tuning);
2842
2843        // Write everything, and churn a tenth of it as we go, which is what
2844        // moves the centroids around under the members that are already filed.
2845        let mut rng = Rng::new(23);
2846        for (i, v) in store.0.iter().enumerate() {
2847            ix.insert(i as u64, v);
2848            if i > 100 && i % 10 == 0 {
2849                let victim = rng.below(i) as u64;
2850                ix.remove(victim);
2851                ix.insert(victim, &store.0[victim as usize]);
2852            }
2853            ix.maintain(&store, 512);
2854        }
2855        ix.maintain(&store, 1 << 20);
2856        consistent(&ix);
2857        assert_eq!(ix.len(), store.0.len());
2858
2859        let r = recall(&ix, &store, 10, 60);
2860        assert!(r >= 0.95, "recall at 10 after the stream was {r}");
2861    }
2862
2863    /// What the sweep is for, measured as the thing it actually fixes rather
2864    /// than through recall.
2865    ///
2866    /// Drift is members filed under a partition that is no longer their nearest,
2867    /// which is what a split leaves behind in the partitions around it. It shows
2868    /// up in recall eventually, but recall is a blunt instrument here and moves
2869    /// by a percent for reasons that have nothing to do with this, so the
2870    /// straight count is the honest measurement.
2871    #[test]
2872    fn the_sweep_is_what_keeps_members_under_their_nearest_centroid() {
2873        let dim = 96;
2874        let store = corpus(dim, 2000, 10, 29);
2875        let tuning = Tuning {
2876            posting: 48,
2877            ..Tuning::default()
2878        };
2879        let with = misfiled(&build(&store, dim, tuning), &store);
2880        let without = misfiled(&build(&store, dim, Tuning { sweep: 0, ..tuning }), &store);
2881        assert!(
2882            with * 4 < without,
2883            "sweeping left {with} members drifted and not sweeping left {without}"
2884        );
2885    }
2886
2887    /// How many members are filed under something that is not their nearest
2888    /// centroid.
2889    /// Asked once per member rather than once per posting entry, because a
2890    /// boundary copy sits in a partition that is not the member's nearest on
2891    /// purpose, and counting one as drift would read replication as the very
2892    /// thing the sweep exists to undo. A member has drifted when none of the
2893    /// partitions holding it is its nearest.
2894    fn misfiled(ix: &Partitions, store: &Store) -> usize {
2895        let mut buf = vec![0.0f32; ix.dim()];
2896        let mut wrong = 0;
2897        for id in 0..store.0.len() as u64 {
2898            if !ix.contains(id) {
2899                continue;
2900            }
2901            assert!(store.get(id, &mut buf));
2902            let near = ix.nearest(&ix.quant.rotate(&buf));
2903            if ix.placed_at(id, near).is_none() {
2904                wrong += 1;
2905            }
2906        }
2907        wrong
2908    }
2909
2910    #[test]
2911    fn a_vector_the_log_forgot_is_dropped_rather_than_returned() {
2912        let dim = 64;
2913        let store = corpus(dim, 400, 4, 31);
2914        let tuning = Tuning {
2915            posting: 24,
2916            ..Tuning::default()
2917        };
2918        let mut ix = build(&store, dim, tuning);
2919        assert!(ix.contains(11));
2920
2921        // The log loses one without telling the index, which is the state a
2922        // crash between two appends leaves behind.
2923        let holey = Holey(store.0.clone(), 11);
2924        assert!(
2925            ix.search(&store.0[11], 5, &holey)
2926                .iter()
2927                .all(|h| h.id != 11)
2928        );
2929
2930        // And maintenance walking over it takes it out for good.
2931        for id in 0..300u64 {
2932            ix.remove(id);
2933        }
2934        ix.maintain(&holey, 1 << 20);
2935        consistent(&ix);
2936        assert!(!ix.contains(11));
2937    }
2938
2939    #[test]
2940    fn rotating_first_is_the_same_as_rotating_inside() {
2941        // The whole index rests on the rotation being linear, so this is the
2942        // property, not an implementation detail.
2943        let dim = 128;
2944        let q = Quantizer::new(dim, Bits::One, 5);
2945        let store = corpus(dim, 2, 1, 37);
2946        let (v, c) = (&store.0[0], &store.0[1]);
2947
2948        let mut a = vec![0u8; q.code_bytes()];
2949        let one = q.encode(v, c, &mut a);
2950        let mut b = vec![0u8; q.code_bytes()];
2951        let two = q.encode_rotated(&q.rotate(v), &q.rotate(c), &mut b);
2952
2953        assert_eq!(a, b, "the two ways round should write the same code");
2954        assert!((one.norm - two.norm).abs() < 1e-4);
2955        assert!((one.scale - two.scale).abs() < 1e-4);
2956    }
2957
2958    #[test]
2959    fn two_means_splits_two_clouds_apart() {
2960        let dim = 8;
2961        let mut xs = Vec::new();
2962        for i in 0..40 {
2963            let far = if i % 2 == 0 { 0.0 } else { 10.0 };
2964            for d in 0..dim {
2965                xs.push(far + (i as f32 + d as f32) * 0.01);
2966            }
2967        }
2968        let (a, b) = two_means(&xs, dim);
2969        let (near, away) = if a[0] < b[0] { (a, b) } else { (b, a) };
2970        assert!(near[0] < 1.0, "{near:?}");
2971        assert!(away[0] > 9.0, "{away:?}");
2972    }
2973}