diskann

package module
v0.0.0-...-7845765 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: May 12, 2026 License: MIT Imports: 10 Imported by: 0

Documentation

Overview

Package diskann implements a pure Go DiskANN (Vamana) index for disk-resident approximate nearest neighbor search on billion-scale datasets.

Key design: graph topology + PQ codes reside in RAM (~100 bytes/vector), while raw embeddings live on disk (mmap) and are fetched only for final reranking. This gives 8-16x RAM reduction compared to in-memory HNSW.

Algorithm reference: Subramanya et al., "DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node", NeurIPS 2019.

Index

Constants

View Source 
const DiskANNMagic uint32 = 0x444B414E // "DKAN"

DiskANNMagic identifies a gleann DiskANN index file.

View Source 
const DiskANNVersion uint32 = 1

DiskANNVersion is the current format version.

Variables

This section is empty.

Functions

func ADCDistance 

func ADCDistance(table [][]float32, codes []byte) float32

ADCDistance computes the approximate distance from a query to a PQ-encoded vector using a precomputed distance table. O(M) operations.

func CosineDistance 

func CosineDistance(a, b []float32) float32

CosineDistance computes 1 - cosine_similarity.

func L2DistanceSquared 

func L2DistanceSquared(a, b []float32) float32

L2DistanceSquared computes the squared Euclidean distance.

Types

type BackendBuilder 

type BackendBuilder interface {
	Build(ctx context.Context, embeddings [][]float32) ([]byte, error)
	AddVectors(ctx context.Context, indexData []byte, embeddings [][]float32, startID int64) ([]byte, error)
	RemoveVectors(ctx context.Context, indexData []byte, ids []int64) ([]byte, error)
}

BackendBuilder constructs and serializes an index.

type BackendSearcher 

type BackendSearcher interface {
	Load(ctx context.Context, indexData []byte, meta IndexMeta) error
	Search(ctx context.Context, query []float32, topK int) ([]int64, []float32, error)
	SearchWithRecompute(ctx context.Context, query []float32, topK int, recompute EmbeddingRecomputer) ([]int64, []float32, error)
	Close() error
}

BackendSearcher loads and queries an index.

type Builder 

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

Builder builds DiskANN indexes.

func (*Builder) AddVectors 

func (b *Builder) AddVectors(ctx context.Context, indexData []byte, embeddings [][]float32, startID int64) ([]byte, error)

func (*Builder) Build 

func (b *Builder) Build(ctx context.Context, embeddings [][]float32) ([]byte, error)

func (*Builder) RemoveVectors 

func (b *Builder) RemoveVectors(ctx context.Context, indexData []byte, ids []int64) ([]byte, error)

type Candidate 

type Candidate struct {
	ID       int64
	Distance float32
}

Candidate is a (id, distance) pair.

func ADCDistanceSorted 

func ADCDistanceSorted(table [][]float32, candidates []int64, codes [][]byte) []Candidate

ADCDistanceSorted computes ADC distances for multiple candidates and returns them sorted. Used for PQ-based candidate prefiltering.

func SearchDiskIndex 

func SearchDiskIndex(idx *DiskIndex, query []float32, topK, searchL, rerankK int, distFunc DistanceFunc, getVector func(int64) []float32) []Candidate

SearchDiskIndex performs DiskANN's beam search with PQ prefiltering.

Algorithm:

  1. Beam search on graph using PQ approximate distances (cheap, RAM-only).
  2. Collect top-L candidates by PQ distance.
  3. Rerank top candidates using exact distances (expensive, disk read).

Parameters:

  • query: the query vector
  • topK: number of results to return
  • searchL: candidate list size (beam width)
  • rerankK: how many PQ candidates to rerank (0 = 2*topK)
  • getVector: callback to fetch raw embedding (from mmap/disk/recompute)

func SearchExact 

func SearchExact(idx *DiskIndex, query []float32, topK, searchL int, distFunc DistanceFunc) []Candidate

SearchExact performs beam search with exact distances (all embeddings in RAM). This is used when all embeddings fit in memory (small datasets or testing).

type Config 

type Config struct {
	IndexDir      string        `json:"index_dir"`
	Backend       string        `json:"backend"`
	DiskANNConfig DiskANNConfig `json:"diskann_config,omitempty"`
}

Config is a minimal standalone config (mirrors gleann.Config pattern).

func DefaultConfig 

func DefaultConfig() Config

DefaultConfig returns a Config with sensible defaults.

type DiskANNConfig 

type DiskANNConfig struct {
	// R is the max out-degree per node (analogous to HNSW M).
	// Default: 64.
	R int `json:"r"`

	// L is the candidate list size during build (analogous to efConstruction).
	// Default: 100.
	L int `json:"l"`

	// Alpha controls the robust prune diversity parameter.
	// α=1.0 = pure distance, α>1.0 = more diversity. Default: 1.2.
	Alpha float64 `json:"alpha"`

	// PQDim is the number of PQ sub-quantizer dimensions.
	// Each sub-quantizer encodes (dims/PQDim) dimensions into 1 byte.
	// Default: 0 (auto = dims/4, clamped to [1, dims]).
	PQDim int `json:"pq_dim"`

	// PQCentroids is the number of centroids per sub-quantizer.
	// Must be ≤ 256 (stored as uint8). Default: 256.
	PQCentroids int `json:"pq_centroids"`

	// SearchL is the candidate list size during search.
	// Default: 100.
	SearchL int `json:"search_l"`

	// SearchPQRerank is how many PQ candidates to rerank with full vectors.
	// Default: 2 * TopK.
	SearchPQRerank int `json:"search_pq_rerank"`

	// DistanceMetric selects distance function. Default: "l2".
	DistanceMetric DistanceMetric `json:"distance_metric,omitempty"`

	// UseMmap enables memory-mapped search for raw vectors.
	UseMmap bool `json:"use_mmap,omitempty"`
}

DiskANNConfig holds build and search parameters.

func (*DiskANNConfig) Defaults 

func (c *DiskANNConfig) Defaults(dims int)

Defaults fills zero-valued fields with sensible defaults.

type DiskIndex 

type DiskIndex struct {
	// Header
	NumNodes int64
	Dims     int
	R        int // max out-degree
	Medoid   int64
	PQM      int // number of sub-quantizers
	PQK      int // centroids per sub-quantizer
	PQSubDim int // sub-vector dimensionality

	// Graph topology: neighbors[i] has neighborOffsets[i]..neighborOffsets[i+1] entries.
	NeighborOffsets []int64
	Neighbors       []int64

	// PQ data
	Codebook *PQCodebook
	PQCodes  [][]byte // PQCodes[nodeIdx] = M-byte code

	// Raw embeddings (all in-memory for Build; mmap'd for search).
	Embeddings [][]float32
}

DiskIndex is the serialized on-disk format for a DiskANN index. It combines: Vamana graph topology + PQ codebook + PQ codes + raw embeddings.

Memory layout when loaded:

  • Graph neighbors (flat arrays): ~R*4 bytes/node (RAM)
  • PQ codes: M bytes/node (RAM)
  • PQ codebook: M*K*SubDim*4 bytes (RAM)
  • Raw embeddings: dims*4 bytes/node (disk/mmap)

Total RAM: ~(R*4 + M) bytes/node. For R=64, M=32: ~288 bytes/node. vs HNSW in-memory: ~(M*2*4 + dims*4) bytes/node. For M=32, dims=128: ~768 bytes/node.

func BuildDiskIndex 

func BuildDiskIndex(embeddings [][]float32, cfg DiskANNConfig) (*DiskIndex, error)

BuildDiskIndex constructs a full DiskANN index from embeddings.

func ReadDiskIndex 

func ReadDiskIndex(r io.Reader) (*DiskIndex, error)

ReadDiskIndex deserializes a DiskIndex from a reader.

func (*DiskIndex) DiskUsage 

func (idx *DiskIndex) DiskUsage() int64

DiskUsage estimates the total disk footprint.

func (*DiskIndex) GetNeighbors 

func (idx *DiskIndex) GetNeighbors(i int64) []int64

GetNeighbors returns the neighbors of node i from the flat arrays.

func (*DiskIndex) RAMUsage 

func (idx *DiskIndex) RAMUsage() int64

RAMUsage estimates the RAM needed when embeddings are on disk.

func (*DiskIndex) WriteTo 

func (idx *DiskIndex) WriteTo(w io.Writer) (int64, error)

WriteTo serializes the DiskIndex to a writer. Format: 

[magic:4][version:4][numNodes:8][dims:4][R:4][medoid:8]
[pqM:4][pqK:4][pqSubDim:4]
[neighborOffsets: (N+1)*8 bytes]
[neighbors: totalEdges*8 bytes]
[PQ centroids: M*K*SubDim*4 bytes]
[PQ codes: N*M bytes]
[raw embeddings: N*dims*4 bytes]

type DistanceFunc 

type DistanceFunc func(a, b []float32) float32

DistanceFunc computes distance between two vectors. Lower values = more similar.

func GetDistanceFunc 

func GetDistanceFunc(metric string) DistanceFunc

GetDistanceFunc returns a DistanceFunc for the given metric name.

type DistanceMetric 

type DistanceMetric string

DistanceMetric selects the distance function. 

const (
	DistanceL2     DistanceMetric = "l2"
	DistanceCosine DistanceMetric = "cosine"
)

type EmbeddingRecomputer 

type EmbeddingRecomputer func(ctx context.Context, ids []int64) ([][]float32, error)

EmbeddingRecomputer retrieves embeddings for given IDs from storage.

type Factory 

type Factory struct{}

Factory creates DiskANN builders and searchers.

func (*Factory) Name 

func (f *Factory) Name() string

func (*Factory) NewBuilder 

func (f *Factory) NewBuilder(config Config) BackendBuilder

func (*Factory) NewSearcher 

func (f *Factory) NewSearcher(config Config) BackendSearcher

type IndexMeta 

type IndexMeta struct {
	Name           string    `json:"name"`
	Backend        string    `json:"backend"`
	EmbeddingModel string    `json:"embedding_model"`
	Dimensions     int       `json:"dimensions"`
	NumPassages    int       `json:"num_passages"`
	CreatedAt      time.Time `json:"created_at"`
	UpdatedAt      time.Time `json:"updated_at"`
	Version        string    `json:"version"`
}

IndexMeta stores metadata about a built index.

type PQCodebook 

type PQCodebook struct {
	// M is the number of sub-quantizers.
	M int
	// K is the number of centroids per sub-quantizer.
	K int
	// SubDim is the dimensionality of each sub-vector.
	SubDim int
	// Dims is the original vector dimensionality (M * SubDim).
	Dims int

	// Centroids[m][k] is the k-th centroid for the m-th sub-quantizer.
	// Shape: [M][K][SubDim]
	Centroids [][][]float32
}

PQCodebook implements Product Quantization for approximate distance computation. It splits D-dimensional vectors into M sub-vectors, each quantized independently using k-means with K centroids. A vector is encoded as M bytes (when K≤256).

Memory: M*K*subDim float32 centroids (~256KB for M=32, K=256, subDim=4). Encoded vector: M bytes (vs D*4 bytes raw).

func TrainPQ 

func TrainPQ(vectors [][]float32, m, k int) (*PQCodebook, error)

TrainPQ trains a PQ codebook from training vectors using k-means.

func (*PQCodebook) BuildDistanceTable 

func (pq *PQCodebook) BuildDistanceTable(query []float32, distFunc DistanceFunc) [][]float32

BuildDistanceTable precomputes distances from a query to all centroids. Returns table[m][k] = distance from query sub-vector m to centroid k. This enables Asymmetric Distance Computation (ADC): O(M) per candidate.

func (*PQCodebook) Encode 

func (pq *PQCodebook) Encode(vec []float32) []byte

Encode encodes a vector into M PQ codes (one byte per sub-quantizer).

func (*PQCodebook) EncodeAll 

func (pq *PQCodebook) EncodeAll(vectors [][]float32) [][]byte

EncodeAll encodes multiple vectors.

type Searcher 

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

Searcher loads and queries a DiskANN index.

func (*Searcher) Close 

func (s *Searcher) Close() error

func (*Searcher) Load 

func (s *Searcher) Load(ctx context.Context, indexData []byte, meta IndexMeta) error

func (*Searcher) Search 

func (s *Searcher) Search(ctx context.Context, query []float32, topK int) ([]int64, []float32, error)

func (*Searcher) SearchWithRecompute 

func (s *Searcher) SearchWithRecompute(ctx context.Context, query []float32, topK int, recompute EmbeddingRecomputer) ([]int64, []float32, error)

type VamanaGraph 

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

VamanaGraph is a single-level navigable graph built with the Vamana algorithm. Unlike HNSW's multi-level structure, Vamana uses one flat graph with robust pruning (α-parameter) for diversity-aware neighbor selection.

func NewVamanaGraph 

func NewVamanaGraph(r, l, dims int, alpha float64, distFunc DistanceFunc) *VamanaGraph

NewVamanaGraph creates an empty Vamana graph.

func (*VamanaGraph) AllNodeIDs 

func (g *VamanaGraph) AllNodeIDs() []int64

AllNodeIDs returns all node IDs in insertion order.

func (*VamanaGraph) Build 

func (g *VamanaGraph) Build(embeddings [][]float32)

Build constructs the Vamana graph from a batch of embeddings. This is the standard two-pass Vamana algorithm: Pass 1: random-order insertion with α=1.0 (greedy) Pass 2: re-insertion with α>1.0 (diverse)

func (*VamanaGraph) GetNeighbors 

func (g *VamanaGraph) GetNeighbors(id int64) []int64

GetNeighbors returns the neighbors of a node.

func (*VamanaGraph) GetVector 

func (g *VamanaGraph) GetVector(id int64) []float32

GetVector returns the vector for a node.

func (*VamanaGraph) Medoid 

func (g *VamanaGraph) Medoid() int64

Medoid returns the entry point ID.

func (*VamanaGraph) Search 

func (g *VamanaGraph) Search(query []float32, topK, l int) []Candidate

Search performs greedy search on the built graph.

func (*VamanaGraph) Size 

func (g *VamanaGraph) Size() int

Size returns the number of nodes.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.