ivfq

README

ivfq

go get github.com/axiomhq/ivfq

Inverted-file clustering plus 1-bit quantization for approximate nearest neighbour search, in pure Go with AVX2 kernels. The FAISS "IVF + quantizer" pair, without the storage: you keep the rows, this package gives you the clusters, the codes and the scan.

package	what it holds
ivfq	sizing rule (HoodK, DefaultNprobe), Policy, split and merge rules, the centroid Tree, the centroid codec, and the distance vocabulary (Metric, Score, Dot, L2Sq, CosineSim)
kmeans	Config.Fit: sampled greedy k-means++ with one parallel assignment pass, two-level above kmeans.TwoLevelK centroids
rabitq	1-bit RaBitQ codes: Quantize, Quantizer, Scorer, Rotation
rank	top-k selection and reciprocal rank fusion
late	late-interaction scoring of multi-vector documents
sparse	the string-keyed sparse vector
codec	f16 and int8 row storage
bench	Big ANN corpus readers and recall@k

No I/O anywhere but bench.

Build an index

k := ivfq.HoodK(len(vectors)) // clusters of about 1,024 rows
centroids, assign, err := kmeans.Config{K: k, Iters: 11, Seed: 1}.Fit(ctx, vectors)
tree, err := ivfq.Build(ctx, centroids, 32)
rot := rabitq.NewRotation(rabitq.Seed("index", "field"), dims)
codes, err := rabitq.Quantize(rowsOfCluster, rabitq.Options{Metric: ivfq.L2, Rotation: rot})

Fit trains on a sample of at most kmeans.SampleSize(k) rows and assigns every row in one parallel pass; the result is bit-identical whatever GOMAXPROCS is. Above kmeans.TwoLevelK (512) it fits √k coarse centroids, then k fine ones under them, and each row searches only its 16 nearest coarse centroids' fine ones: about 17·√k distances per row, not k. To route rows later, s := tree.NewSearcher() then s.Assign(row, beam), one Searcher per goroutine. One bit per dimension plus two float32 per row.

Query

probe := tree.Nearest(query, ivfq.DefaultNprobe(k))
sc := codes.Scorer(rabitq.NewQuery(query, ivfq.L2, rot))
score, bound := sc.ScoreAndBound(row)   // estimate and upper bound
exact := ivfq.Score(ivfq.L2, query, raw) // rerank the shortlist

A row may be skipped by an exact top-k cutoff only when its bound is below the cutoff. The bound is probabilistic; Query.Exact drops it.

Splits and merges keep the clusters at their size as rows arrive: a split is kmeans.Config{K: 2, Iters: kmeans.SplitIters}, and Policy, MergeTargets, SplitAbove and MergeBelow decide when, reading only per-cluster counts and radii. No rule ever touches every row.

Rank

rank orders hits by score descending, ties by id ascending.

1. Select the top k as you score: sel := rank.NewTopK(k), sel.Push(id, score) per candidate, then sel.Hits(). rank.Select(hits, k) does the same for a slice.
2. Fuse ranked legs: rank.RRF(k, legs...), or rank.Fusion{RankConstant: c, Weights: w}.RRF(k, legs...) after Validate(len(legs)).
3. Late interaction: late.Score(metric, queryTokens, docTokens) sums, over query tokens, the best document-token score. Check inputs with late.Validate first.

Benchmark data

bench reads the Big ANN corpora and scores recall. Stdlib only.

1. Read a corpus: b, err := bench.OpenBin(path), then b.Next(n) in batches for .u8bin/.fbin.
2. Read the ground truth: bench.ReadGroundTruth(path) for .ibin.
3. Score: bench.RecallAtK(got, truth, 10).
4. CPU time: bench.CPUSeconds() before and after a run.

Kernels

Every SIMD kernel lives in internal/simd with a pure Go reference and a test that pins it. amd64 picks the kernel at run time from golang.org/x/sys/cpu; every other architecture, and amd64 without the flags, runs the Go reference.

kernel	amd64 path	needs	claim	test
Dots	dotsAVX2FMA	AVX2 + FMA	each out[r] equals Dot(q, row r) bit for bit	TestDotsMatchesDot
1-bit scan	bitProductAVX2	AVX2 + POPCNT	popcounts identical to the Go loop	TestBitProductMatchesGeneric
rotation	rotatePairsAVX2	AVX2	no FMA, so equal to the Go loop bit for bit	TestRotatePairsMatchesGeneric
i8 codec	decodeI8AVX2, encodeI8AVX2	AVX2	same bytes and floats as the scalar codec	TestI8CodecsMatchScalar

Dot, L2Sq and CosineSim keep several float32 accumulator lanes, so they can differ from a float64 loop by about 1e-7 relative. Rankings do not change, but centroids are only reproducible against a fixed kernel: treat a kernel change like a seed change.

Benchmark

BenchmarkBitScore, 10,000 rows of 128 dimensions, AMD EPYC 7502P (AVX2), Go 1.27, n=10:

	per 10k rows	per row
1-bit scan, estimate and bound	427.6 µs ± 1%	43 ns
f32 L2Sq over raw rows	258.9 µs ± 2%	26 ns

Test

go test -race ./...
go test -run '^$' -bench BitScore -count 10 ./rabitq

License

MIT, see LICENSE.

Documentation

Overview

Package ivfq is the clustering half of an IVF index in the FAISS sense ("IVF + quantizer"): the sizing rule that keeps clusters at about HoodTarget vectors, the centroid Tree a query probes, the split and merge rules that hold cluster sizes between rebuilds, the Policy that names the next cluster to split, and a binary codec for centroid sets and their deltas. It also holds the distance vocabulary every subpackage shares: Metric, Score, Dot, L2Sq and CosineSim.

The other halves are subpackages: kmeans fits the centroids, rabitq gives every cluster a 1-bit candidate column with a scorer and an error bound, rank selects and fuses hits, late scores multi-vector documents, sparse is the string-keyed sparse vector, codec stores rows at f16 and int8, and bench reads the standard ANN corpora. No I/O anywhere but bench.

Index

• Constants
• Variables
• func ApplyCentroidDeltas(base [][]float32, deltas ...[]CentroidUpsert) ([][]float32, error)
• func Centroid(vecs [][]float32, dims int, spherical bool) ([]float32, float64)
• func CosineSim(a, b []float32) float32
• func DecodeCentroids(data []byte, dims int) ([][]float32, error)
• func DefaultNprobe(k int) int
• func Dot(a, b []float32) float32
• func Dots(q, block, out []float32)
• func EncodeCentroidDelta(entries []CentroidUpsert) []byte
• func EncodeCentroids(cs [][]float32) []byte
• func HoodK(total int) int
• func HoodSize(total int) int
• func L2Sq(a, b []float32) float32
• func MeanDistance(vecs [][]float32, c []float32) float64
• func MergeBelow(total int) int
• func MergeOwed(stats []ClusterStat) bool
• func MergeTargets(stats []ClusterStat, centroids [][]float32, busy func(i int) bool, n int) [][2]int
• func Normalize(v []float32)
• func RoundSqrt(total int) int
• func Score(metric Metric, q, v []float32) float32
• func SplitAbove(total int) int
• func TotalCount(stats []ClusterStat) int
• type CentroidUpsert
  • func CentroidDelta(prev, next [][]float32, limit int) ([]CentroidUpsert, bool)
  • func DecodeCentroidDelta(data []byte, dims int) ([]CentroidUpsert, error)
• type ClusterStat
• type DefaultPolicy
  • func (DefaultPolicy) Bootstrap(total int) (int, bool)
  • func (DefaultPolicy) SplitTarget(clusters []ClusterStat) (int, bool)
• type Metric
  • func ParseMetric(s string) (Metric, error)
  • func (m Metric) MarshalText() ([]byte, error)
  • func (m Metric) String() string
  • func (m *Metric) UnmarshalText(text []byte) error
• type Policy
• type Searcher
  • func (s *Searcher) Assign(q []float32, beam int) int
  • func (s *Searcher) Nearest(q []float32, probe int) []int
  • func (s *Searcher) TwoStageAssign(q []float32, top int) int
• type Tree
  • func Build(ctx context.Context, centroids [][]float32, fanout int) (Tree, error)
  • func UnmarshalTree(data []byte) (Tree, error)
  • func (t *Tree) ApplyCentroidDeltas(deltas [][]CentroidUpsert) error
  • func (t Tree) Clone() Tree
  • func (t Tree) Dims() int
  • func (t Tree) Evaluations(q []float32, probe int) ([]int, int)
  • func (t Tree) EvaluationsBeam(q []float32, beam, probe int) ([]int, int)
  • func (t Tree) EvaluationsBeamI8(q []float32, beam, probe int) ([]int, int)
  • func (t Tree) Height() int
  • func (t Tree) Leaf(id int) []float32
  • func (t Tree) Leaves() int
  • func (t Tree) MarshalBinary() ([]byte, error)
  • func (t Tree) Nearest(q []float32, probe int) []int
  • func (t *Tree) NewSearcher() *Searcher
  • func (t *Tree) QuantizeI8()
  • func (t *Tree) Truncate(n int) error
  • func (t Tree) TwoStageNearest(q []float32, top int) ([]int, int)
  • func (t *Tree) Upsert(id int, v []float32) error
  • func (t *Tree) UpsertAll(centroids [][]float32, limit int) (*Tree, bool)

Constants

const (
	MergeMax = 3

	// HoodTarget is the vectors-per-neighborhood the sizing rule aims at.
	// It is the ONE constant that makes a query's cost constant in N: three
	// things are proportional to it and they pull opposite ways — candidates
	// scored per query is nprobe x HoodTarget, pack bytes per hood is
	// HoodTarget x dims x 4 (512 KiB at 128 dims, one object read's worth
	// and the unit a cache evicts), and the centroid set both the query and
	// the build must scan is N/HoodTarget.
	//
	// 1024 was CHOSEN, not swept: it makes the 1M index 977 hoods against the
	// 1,000 the old rule gave, so the 1M measurements are a controlled A/B on
	// one corpus rather than a comparison of two different indexes. A sweep
	// over {512, 2048, 4096} at 1M is the obvious follow-up and has not been
	// run. The reasoning: fixed-size neighborhoods.
	HoodTarget = 1024

	// CandidateBudget is how many candidates a default query is allowed to
	// score. 100,000 is what the sqrt-sized index spent at 1M for recall
	// 0.968, so the default is the old default, made constant in N.
	CandidateBudget = 100_000
)

Variables

var ErrCorrupt = errors.New("ivfq: corrupt")

ErrCorrupt wraps every decode failure of a centroid set or delta.

Functions

func ApplyCentroidDeltas

func ApplyCentroidDeltas(base [][]float32, deltas ...[]CentroidUpsert) ([][]float32, error)

ApplyCentroidDeltas folds deltas, oldest first, over a base set. The base slice is not mutated; its vectors are shared. A slot past the end of the set (a gap) is corruption: appends are contiguous by construction.

func Centroid

func Centroid(vecs [][]float32, dims int, spherical bool) ([]float32, float64)

Centroid is the mean of the dims-wide vectors (spherical: normalized) and their mean distance to it; nil for none. Vectors of another width are skipped.

func CosineSim

func CosineSim(a, b []float32) float32

CosineSim returns cosine similarity (higher = more similar). Zero vectors score 0.

func DecodeCentroids

func DecodeCentroids(data []byte, dims int) ([][]float32, error)

DecodeCentroids reads what EncodeCentroids wrote, every centroid dims wide.

func DefaultNprobe

func DefaultNprobe(k int) int

DefaultNprobe is how many hoods a query probes when the caller does not name a number and no recall controller has tuned the index. It is also where package recall's controller starts and what its ceiling is a multiple of, so it is a starting point rather than the last word. Two bounds, and the smaller wins:

max(budgetNprobe, k/32)  the FLOOR: a constant candidate budget, raised
                         to a fixed fraction of the index once that is
                         bigger, because a constant budget does not hold
                         recall (see recallProbeDiv).
max(16, k/4)             the CEILING: never read more than a quarter of
                         the hoods by default once there are enough of
                         them to matter.

The ceiling exists because the budget alone made the default an EXACT scan just below its own size: a 100,000-vector index has 98 hoods and the budget asks for 98, so the default read every one — 4.42 ms measured, for recall the index did not need. The quarter it now reads there is 24 hoods: recall@10 1.000 at 1.12 ms, 24,641 candidates against 100,000. Above ~400,000 vectors the ceiling is never the binding bound and 1M and 10M are untouched. Below 64 hoods it floors at 16, which is at or above k for every index the bootstrap threshold allows, so a small index still probes everything and answers exactly.

func Dot

func Dot(a, b []float32) float32

Dot returns the inner product. With precomputed squared norms it gives a squared distance without the sqrt L2Sq pays: |a-b|^2 = |a|^2 + |b|^2 - 2 a.b. The subtraction loses a few low bits when a and b are much longer than their difference, so callers rank with it; they do not report it.

func Dots

func Dots(q, block, out []float32)

Dots sets out[r] to Dot(q, block[r*len(q):(r+1)*len(q)]) for every row of block, bit for bit, in one call.

func EncodeCentroidDelta

func EncodeCentroidDelta(entries []CentroidUpsert) []byte

EncodeCentroidDelta writes upserts as: magic, dims, count, then per entry the slot and the vector. The dims are carried so a delta applied to the wrong field fails on its header, not on a misaligned float.

func EncodeCentroids

func EncodeCentroids(cs [][]float32) []byte

EncodeCentroids writes a little-endian uint32 count, then each centroid's float32 values.

func HoodK

func HoodK(total int) int

HoodK is the sizing rule: k = ceil(total / HoodTarget), so hood size is constant in N and hood COUNT carries the growth. It replaces round(sqrt(total)), under which both count and size grew as sqrt(N) and candidates per query therefore grew linearly in N. Identical to the old rule for total <= 256; larger for everything above, by design.

func HoodSize

func HoodSize(total int) int

HoodSize is the per-hood target HoodK implies at this total: HoodTarget once the index is big enough to want more than hoodFloor hoods, and total/k below that. Split and merge are stated against THIS, not against HoodTarget directly, so the small-index floor does not make every hood look tiny and merge the index into one pack.

func L2Sq

func L2Sq(a, b []float32) float32

L2Sq returns squared euclidean distance. Squared preserves ranking and skips the sqrt.

func MeanDistance

func MeanDistance(vecs [][]float32, c []float32) float64

MeanDistance is the mean Euclidean distance from c to the vectors as wide as c; 0 for none.

func MergeBelow

func MergeBelow(total int) int

MergeBelow is the hood size under which a hood is a wasted probe slot and should be absorbed by its nearest neighbor. Never below MergeMax, so the pre-existing "singleton hoods merge" behavior survives at any size.

func MergeOwed

func MergeOwed(stats []ClusterStat) bool

MergeOwed reports an empty cluster (reclaim) or a tiny one with room somewhere to merge into. MergeTargets refuses a merge past the split trigger, so a tiny no cluster has room for owes nothing: counting it would buy a rebalancing cycle that declines every time.

func MergeTargets

func MergeTargets(stats []ClusterStat, centroids [][]float32, busy func(i int) bool, n int) [][2]int

MergeTargets plans up to n merges at once: pairs of a tiny cluster and the nearest cluster to absorb it, the smallest tinies first, each tiny merging into its nearest neighbour when that neighbour stays at or under the split trigger with its rows. A merge past the trigger undoes the split that made the tiny, and the next cycle splits the host again: a lopsided 2-means split leaves a child under MergeBelow, the child's nearest cluster is its sibling, and the layout alternates split and merge on every cycle, forever. A tiny with no such neighbour stays until it grows. The size trigger is the one mirrored: a merge can still widen a host past the radius trigger.

No cluster is in two pairs, and none busy reports (nil: none) is in any: a tiny whose nearest neighbour is busy or already paired waits for a later call rather than taking its second nearest. centroids[i] is stats[i]'s centroid; pairs index stats as {tiny, host}. At most mergeTries distance scans (k × dims each) per pair wanted; they run on GOMAXPROCS goroutines, a batch of tinies per pass over the centroids.

func Normalize

func Normalize(v []float32)

Normalize scales v to unit length in place. A zero-norm v is left as is.

func RoundSqrt

func RoundSqrt(total int) int

RoundSqrt is the k heuristic: about √N neighborhoods of about √N vectors.

func Score

func Score(metric Metric, q, v []float32) float32

Score is q against v under metric as a "higher is better" score: the negated squared distance for L2, the similarity for Cosine, the inner product for InnerProduct. An unknown Metric is a programming error and panics.

func SplitAbove

func SplitAbove(total int) int

SplitAbove is the band's MAXIMUM: the hood size past which the policy owes a split (SplitTarget's size trigger, 1.4x the target). It is stated here rather than inlined at the trigger because a rewrite that leaves a hood above this may split it inside the rewrite, the way SPFresh's LIRE splits a posting inside the insert that overflowed it, and both readers of the rule have to agree on the same number. 1,432 at a million vectors.

func TotalCount

func TotalCount(stats []ClusterStat) int

TotalCount sums a stats slice — the live vector count the split and merge thresholds are stated against. Both triggers are relative to the size the sizing rule wants at THIS total, so neither needs the policy to be told N.

Types

type CentroidUpsert

type CentroidUpsert struct {
	ID  int
	Vec []float32
}

CentroidUpsert is one entry of a centroid delta: the vector now at slot ID, which is either an existing slot (a recentred cluster) or one past the last (a cluster a split added), or, with a nil Vec, the set cut to its first ID slots (the last entry of a delta, after its upserts). A caller that drops a cluster moves the last one into its slot and cuts the end, so a merge is two entries and not a full set.

func CentroidDelta

func CentroidDelta(prev, next [][]float32, limit int) ([]CentroidUpsert, bool)

CentroidDelta reports the upserts that turn prev into next when next keeps every slot of prev in place and only replaces or appends at most limit slots. Callers pass max(2, len(prev)/16): small enough that a delta stays cheap to apply on every load, and at least two so a split always qualifies. ok=false means write a full set: slots moved or vanished, or more than limit changed.

func DecodeCentroidDelta

func DecodeCentroidDelta(data []byte, dims int) ([]CentroidUpsert, error)

DecodeCentroidDelta reads what EncodeCentroidDelta wrote. Like DecodeCentroids it grows with bytes actually read, so a lying count dies on its first truncated entry instead of sizing an allocation.

type ClusterStat

type ClusterStat struct {
	ID, Count int
	Radius    float64
	Hits      int
}

ClusterStat is one hood as its index records it. Count is the number of vectors the hood holds and Radius the mean distance from those members to its centroid, both measured by whatever last rewrote the hood, so both are exact and free to read. Radius == 0 means either "every member is the same point" or "emptied, awaiting reclaim"; either way it is not splittable evidence. Hits is query heat since the last cycle, and is 0 on a detached cycle.

type DefaultPolicy

type DefaultPolicy struct{}

DefaultPolicy is the one shipped policy; the Policy interface exists so a test can substitute a stub that names a split unconditionally. SplitTarget answers with one hood, chosen by size and then by radius.

func (DefaultPolicy) Bootstrap

func (DefaultPolicy) Bootstrap(total int) (int, bool)

func (DefaultPolicy) SplitTarget

func (DefaultPolicy) SplitTarget(clusters []ClusterStat) (int, bool)

type Metric

type Metric int

Metric is the distance a field is scored with. The zero value is L2.

const (
	L2           Metric = iota // squared euclidean distance, negated by Score
	Cosine                     // cosine similarity
	InnerProduct               // inner product
)

func ParseMetric

func ParseMetric(s string) (Metric, error)

ParseMetric is the inverse of Metric.String: "l2", "cosine" or "ip".

func (Metric) MarshalText

func (m Metric) MarshalText() ([]byte, error)

MarshalText implements encoding.TextMarshaler, so a Metric reads and writes as its name in JSON and other text configs.

func (Metric) String

func (m Metric) String() string

func (*Metric) UnmarshalText

func (m *Metric) UnmarshalText(text []byte) error

UnmarshalText implements encoding.TextUnmarshaler.

type Policy

type Policy interface {
	// Bootstrap: no index exists and total is the live vector count.
	// ok=false leaves the index unbuilt this cycle. The build that follows
	// is the ONE pass over all N vectors the design allows.
	Bootstrap(total int) (k int, ok bool)
	// SplitTarget: consulted AFTER the append pass with post-append counts.
	SplitTarget(clusters []ClusterStat) (id int, ok bool)
}

Policy implementations must be safe for concurrent use; the shipped policies are stateless.

Every trigger past Bootstrap is stated on []ClusterStat and nothing else, and that is the point: per-cluster counts and radii are what an index already holds, so consulting the policy costs O(k) arithmetic and no pass over the vectors. There is deliberately NO whole-index trigger: nothing a Policy answers can ask for a pass over every vector. Split, merge and reclaim are the whole vocabulary, and rebuilding an index is a caller's explicit decision.

type Searcher

type Searcher struct {
	// contains filtered or unexported fields
}

Searcher is Nearest, Assign and TwoStageAssign with scratch that survives between calls: a warm Searcher allocates nothing per query. It belongs to one goroutine; a bulk caller makes one per worker. Nearest allocates its candidate buffer per call, about 45 KB per row when a re-route runs every member of a cluster through it.

func (*Searcher) Assign

func (s *Searcher) Assign(q []float32, beam int) int

Assign is the leaf a beam-wide search ranks first, or -1 when the tree is empty or q is the wrong width: the same leaf Nearest(q, beam)[0] names.

func (*Searcher) Nearest

func (s *Searcher) Nearest(q []float32, probe int) []int

Nearest is Tree.Nearest: up to probe leaf ids, nearest first, in the same order. The slice is the Searcher's own and the next call overwrites it.

func (*Searcher) TwoStageAssign

func (s *Searcher) TwoStageAssign(q []float32, top int) int

TwoStageAssign ranks one leaf in two exact stages: a distance for every child of the root, then a distance for every leaf under the top subtrees. It is TwoStageNearest(q, top)[0], or -1.

Measured, it is a wash for bulk routing: on real BIGANN rows the stage-two scan is the whole cost and the beam's pruning wins on large trees, so the maintainer's router keeps the beam. Small trees are where this wins: at k = 977 a top-4 two-stage reached 99.1% agreement at 2.5x less wall per row than beam 16, because the stage-one fan there is 10 wide and four subtrees cover nearly the whole set. Choose by the benchmark, not by the evaluation count.

type Tree

type Tree struct {
	// contains filtered or unexported fields
}

Tree is a deterministic hierarchy over a set of leaf centroids.

func Build

func Build(ctx context.Context, centroids [][]float32, fanout int) (Tree, error)

Build constructs a fanout-ary hierarchy. Its fixed seeds make identical inputs deterministic; fanout values below two are promoted to two. A cancelled ctx returns its error and no tree.

func UnmarshalTree

func UnmarshalTree(data []byte) (Tree, error)

UnmarshalTree decodes MarshalBinary output. Counts are checked against the bytes remaining before they size an allocation, and indices against the encoder's invariants: in range, the root is node zero, every child numbered after its parent (Build appends a node before recursing into its children, which is what makes traversal finite), every other node referenced by exactly one parent, every leaf owned by one node, and no node both branches and holds leaves.

func (*Tree) ApplyCentroidDeltas

func (t *Tree) ApplyCentroidDeltas(deltas [][]CentroidUpsert) error

ApplyCentroidDeltas upserts deltas, oldest first, into t in place.

func (Tree) Clone

func (t Tree) Clone() Tree

Clone returns a deep copy of the tree: every centroid, child and leaf list, scan block, and norm slice is duplicated, and the clone's leaves address the clone's blocks. Neither tree can observe the other's Upsert: a decoded tree sits in a shared cache behind concurrent searches, so a caller that means to mutate clones first and swaps the result in.

func (Tree) Dims

func (t Tree) Dims() int

Dims is the width of every centroid in the tree.

func (Tree) Evaluations

func (t Tree) Evaluations(q []float32, probe int) ([]int, int)

Evaluations is Nearest with the number of centroid distance evaluations.

func (Tree) EvaluationsBeam

func (t Tree) EvaluationsBeam(q []float32, beam, probe int) ([]int, int)

EvaluationsBeam is Evaluations with the beam set apart from probe: the search keeps beam nodes per level and ranks every leaf under them, so a beam of probe scans every leaf once probe exceeds the root's fanout (all 5,464 at 1M rows in hoods of 256, for 96 probes). Leaves are ranked exactly; a narrower beam can only miss a leaf under a node it dropped. beam below probe is raised to... nothing: it may be smaller than probe, and the result is still the probe nearest leaves the beam reached.

func (Tree) EvaluationsBeamI8

func (t Tree) EvaluationsBeamI8(q []float32, beam, probe int) ([]int, int)

EvaluationsBeamI8 is EvaluationsBeam, the same leaves in the same order, with the leaf scan read through the holders' int8 copies (QuantizeI8). The beam over internal nodes is EvaluationsBeam's own, in float32, so it visits the same holders. Each scanned leaf gets an int8 distance and an exact bound on its error,

|<q,c> - <q~,c~>| <= ||q|| ||c - c~|| + ||q - q~|| ||c~||

and its float32 distance is computed, as EvaluationsBeam computes it, only for the probe nearest by int8 and for any other leaf whose lower bound is within the probe-th float32 distance so far. The count is the float32 distances computed, internal nodes included. A tree without its int8 copies (not quantized, or changed since) is searched by EvaluationsBeam.

func (Tree) Height

func (t Tree) Height() int

Height returns the number of node levels, including the root and leaves.

func (Tree) Leaf

func (t Tree) Leaf(id int) []float32

Leaf returns leaf id's centroid, aliasing the tree's own storage: read only.

func (Tree) Leaves

func (t Tree) Leaves() int

Leaves is the number of leaf centroids the tree ranks over.

func (Tree) MarshalBinary

func (t Tree) MarshalBinary() ([]byte, error)

MarshalBinary encodes the hierarchy and leaf vectors in a versioned format.

func (Tree) Nearest

func (t Tree) Nearest(q []float32, probe int) []int

Nearest returns up to probe leaf centroid indices, nearest first.

func (*Tree) NewSearcher

func (t *Tree) NewSearcher() *Searcher

NewSearcher returns a Searcher over t. t must not change while the Searcher is in use.

func (*Tree) QuantizeI8

func (t *Tree) QuantizeI8()

QuantizeI8 gives every leaf holder an int8 copy of its leaves for EvaluationsBeamI8: a quarter of the bytes the leaf scan reads, which at 768 dimensions was 12.7 MB a query. It is derived, never serialized, and dropped by any change to the tree (Upsert, Truncate), after which EvaluationsBeamI8 reads the float32 leaves until QuantizeI8 runs again. It writes t, so call it before t is shared.

func (*Tree) Truncate

func (t *Tree) Truncate(n int) error

Truncate cuts the tree to its first n leaves, each later leaf detached from its holder. Internal centroids stay, as they do under Upsert, and a holder may be left empty; the next Build rebalances.

func (Tree) TwoStageNearest

func (t Tree) TwoStageNearest(q []float32, top int) ([]int, int)

TwoStageNearest is TwoStageAssign's search with the picked leaf ids, nearest first, and the number of centroid distance evaluations returned — the pairing Nearest and Evaluations give the beam search.

func (*Tree) Upsert

func (t *Tree) Upsert(id int, v []float32) error

Upsert replaces leaf id's centroid with v, or appends v as a new leaf when id equals the number of leaves; any other id is an error, since appends must stay contiguous. The height never changes and the result round-trips through MarshalBinary and UnmarshalTree.

A replace moves the leaf to the holder nearest to v (replace), so a slot refilled with a far centroid stays reachable. Internal centroids are deliberately not recomputed: the beam walks a handful of candidates per level, and holders are ranked by their means, which drift negligibly over a run of upserts. Recentring is what the next full Build is for.

An append descends the hierarchy the way Nearest does (beam = fanout) to find the holder whose centroid is nearest to v, appends the leaf there, and gives that holder a fresh block — holder blocks are sub-slices of the one leaf arena, so growing in place would trample a neighbouring hood — while re-pointing the holder's leaves at the new block. A holder is allowed to grow past the fanout: search scans every leaf of a visited holder, so correctness holds and quality degrades slowly; rebuild once enough upserts have accumulated.

func (*Tree) UpsertAll

func (t *Tree) UpsertAll(centroids [][]float32, limit int) (*Tree, bool)

UpsertAll turns t into centroids' tree without a Build: t itself when they already agree, and otherwise a clone with the changed and appended slots upserted — what ApplyCentroidDeltas does with a delta. A split changes one slot and appends one, and rebuilding a ~10k-centroid tree for that costs far more than two upserts. A shorter set cuts the tree to its length first (Truncate). At most limit slots are upserted (callers pass max(2, t.Leaves()/16), the CentroidDelta bound); past it ok is false and the caller builds.