Struct MmapIndex 

pub struct MmapIndex {
    pub path: String,
    pub metadata: Metadata,
    pub codec: ResidualCodec,
    pub ivf: Array1<i64>,
    pub ivf_lengths: Array1<i32>,
    pub ivf_offsets: Array1<i64>,
    pub doc_lengths: Array1<i64>,
    pub doc_offsets: Array1<usize>,
    pub mmap_codes: MmapNpyArray1I64,
    pub mmap_residuals: MmapNpyArray2U8,
}

A memory-mapped PLAID index for multi-vector search.

This struct uses memory-mapped files for the large arrays (codes and residuals) instead of loading them entirely into RAM. Only small tensors (centroids, bucket weights, IVF) are loaded into memory.

Memory Usage

Only small tensors (~50 MB for SciFact 5K docs) are loaded into RAM, with code and residual data accessed via OS-managed memory mapping.

Usage

use next_plaid::MmapIndex;

let index = MmapIndex::load("/path/to/index")?;
let results = index.search(&query, &params, None)?;

Fields

path: String

Path to the index directory

metadata: Metadata

Index metadata

codec: ResidualCodec

Residual codec for quantization/decompression

ivf: Array1<i64>

IVF data (concatenated passage IDs per centroid)

ivf_lengths: Array1<i32>

IVF lengths (number of passages per centroid)

ivf_offsets: Array1<i64>

IVF offsets (cumulative offsets into ivf array)

doc_lengths: Array1<i64>

Document lengths (number of tokens per document)

doc_offsets: Array1<usize>

Cumulative document offsets for indexing into codes/residuals

mmap_codes: MmapNpyArray1I64

Memory-mapped codes array (public for search access)

mmap_residuals: MmapNpyArray2U8

Memory-mapped residuals array (public for search access)

Implementations

impl MmapIndex

pub fn load(index_path: &str) -> Result<Self>

Load a memory-mapped index from disk.

This creates merged files for codes and residuals if they don’t exist, then memory-maps them for efficient access.

pub fn get_candidates(&self, centroid_indices: &[usize]) -> Vec<i64>

Get candidate documents from IVF for given centroid indices.

pub fn get_document_embeddings(&self, doc_id: usize) -> Result<Array2<f32>>

Get document embeddings by decompressing codes and residuals.

pub fn get_document_codes(&self, doc_ids: &[usize]) -> Vec<Vec<i64>>

Get codes for a batch of document IDs (for approximate scoring).

pub fn decompress_documents( &self, doc_ids: &[usize], ) -> Result<(Array2<f32>, Vec<usize>)>

Decompress embeddings for a batch of document IDs.

pub fn search( &self, query: &Array2<f32>, params: &SearchParameters, subset: Option<&[i64]>, ) -> Result<SearchResult>

Search for similar documents.

Arguments

• query - Query embedding matrix [num_tokens, dim]
• params - Search parameters
• subset - Optional subset of document IDs to search within

Returns

Search result containing top-k document IDs and scores.

pub fn search_batch( &self, queries: &[Array2<f32>], params: &SearchParameters, parallel: bool, subset: Option<&[i64]>, ) -> Result<Vec<SearchResult>>

Search for multiple queries in batch.

Arguments

• queries - Slice of query embedding matrices
• params - Search parameters
• parallel - If true, process queries in parallel using rayon
• subset - Optional subset of document IDs to search within

Returns

Vector of search results, one per query.

pub fn num_documents(&self) -> usize

Get the number of documents in the index.

pub fn num_embeddings(&self) -> usize

Get the total number of embeddings in the index.

pub fn num_partitions(&self) -> usize

Get the number of partitions (centroids).

pub fn avg_doclen(&self) -> f64

Get the average document length.

pub fn embedding_dim(&self) -> usize

Get the embedding dimension.

pub fn reconstruct(&self, doc_ids: &[i64]) -> Result<Vec<Array2<f32>>>

Reconstruct embeddings for specific documents.

This method retrieves the compressed codes and residuals for each document from memory-mapped files and decompresses them to recover the original embeddings.

Arguments

• doc_ids - Slice of document IDs to reconstruct (0-indexed)

Returns

A vector of 2D arrays, one per document. Each array has shape [num_tokens, dim].

Example

use next_plaid::MmapIndex;

let index = MmapIndex::load("/path/to/index")?;
let embeddings = index.reconstruct(&[0, 1, 2])?;

for (i, emb) in embeddings.iter().enumerate() {
    println!("Document {}: {} tokens x {} dim", i, emb.nrows(), emb.ncols());
}

pub fn reconstruct_single(&self, doc_id: i64) -> Result<Array2<f32>>

Reconstruct a single document’s embeddings.

Convenience method for reconstructing a single document.

Arguments

• doc_id - Document ID to reconstruct (0-indexed)

Returns

A 2D array with shape [num_tokens, dim].

pub fn create_with_kmeans( embeddings: &[Array2<f32>], index_path: &str, config: &IndexConfig, ) -> Result<Self>

Create a new index from document embeddings with automatic centroid computation.

This method:

1. Computes centroids using K-means
2. Creates index files on disk
3. Loads the index using memory-mapped I/O

Note: During creation, data is temporarily held in RAM for processing, then written to disk and loaded as mmap.

Arguments

• embeddings - List of document embeddings, each of shape [num_tokens, dim]
• index_path - Directory to save the index
• config - Index configuration

Returns

The created MmapIndex

pub fn update( &mut self, embeddings: &[Array2<f32>], config: &UpdateConfig, ) -> Result<Vec<i64>>

Update the index with new documents, matching fast-plaid behavior.

This method adds new documents to an existing index with three possible paths:

1. Start-from-scratch mode (num_documents <= start_from_scratch):

   • Loads existing embeddings from embeddings.npy if available
   • Combines with new embeddings
   • Rebuilds the entire index from scratch with fresh K-means
   • Clears embeddings.npy if total exceeds threshold

2. Buffer mode (total_new < buffer_size):

   • Adds new documents to the index without centroid expansion
   • Saves embeddings to buffer for later centroid expansion

3. Centroid expansion mode (total_new >= buffer_size):

   • Deletes previously buffered documents
   • Expands centroids with outliers from combined buffer + new embeddings
   • Re-indexes all combined embeddings with expanded centroids

Arguments

• embeddings - New document embeddings to add
• config - Update configuration

Returns

Vector of document IDs assigned to the new embeddings

pub fn update_with_metadata( &mut self, embeddings: &[Array2<f32>], config: &UpdateConfig, metadata: Option<&[Value]>, ) -> Result<Vec<i64>>

Update the index with new documents and optional metadata.

Arguments

• embeddings - New document embeddings to add
• config - Update configuration
• metadata - Optional metadata for new documents

Returns

Vector of document IDs assigned to the new embeddings

pub fn update_or_create( embeddings: &[Array2<f32>], index_path: &str, index_config: &IndexConfig, update_config: &UpdateConfig, ) -> Result<(Self, Vec<i64>)>

Update an existing index or create a new one if it doesn’t exist.

Arguments

• embeddings - Document embeddings to add
• index_path - Directory for the index
• index_config - Configuration for index creation
• update_config - Configuration for updates

Returns

A tuple of (MmapIndex, Vec<i64>) containing the index and document IDs

pub fn delete(&mut self, doc_ids: &[i64]) -> Result<usize>

Delete documents from the index.

Arguments

• doc_ids - Slice of document IDs to delete (0-indexed)

Returns

The number of documents actually deleted

pub fn delete_with_options( &mut self, doc_ids: &[i64], delete_metadata: bool, ) -> Result<usize>

Delete documents from the index with control over metadata deletion.

Arguments

• doc_ids - Slice of document IDs to delete
• delete_metadata - If true, also delete from metadata.db if it exists

Returns

The number of documents actually deleted