Struct ProductQuantizer 

pub struct ProductQuantizer {
    pub num_subvectors: usize,
    pub num_centroids: usize,
    pub dim: usize,
    pub sub_dim: usize,
    pub codebook: Vec<f32>,
}

Product Quantizer with learned codebooks.

Fields

num_subvectors: usize

Number of sub-vector segments.

num_centroids: usize

Number of centroids per sub-vector (max 256 for u8 codes).

dim: usize

Original vector dimension.

sub_dim: usize

Dimension of each sub-vector.

codebook: Vec<f32>

Codebook: [num_subvectors][num_centroids][sub_dim] stored flat. Layout: codebook[sv * num_centroids * sub_dim + c * sub_dim + d]

Implementations

impl ProductQuantizer

pub fn train( vectors: &[Vec<f32>], dim: usize, num_subvectors: usize, num_centroids: usize, ) -> Self

Train a product quantizer from a set of vectors using k-means.

num_subvectors must evenly divide the vector dimension. num_centroids must be <= 256 (for u8 encoding).

pub fn encode(&self, vector: &[f32]) -> Vec<u8>

Encode a vector into PQ codes (one u8 per subvector).

pub fn decode(&self, codes: &[u8]) -> Vec<f32>

Decode PQ codes back to an approximate vector.

pub fn precompute_distance_table(&self, query: &[f32]) -> Vec<f32>

Precompute distance table for asymmetric distance computation.

Returns a table of shape [num_subvectors][num_centroids] (stored flat) containing the squared L2 distance from each query sub-vector to each centroid.

pub fn asymmetric_distance_with_table(&self, table: &[f32], codes: &[u8]) -> f32

Compute asymmetric distance between query and encoded vector.

Uses a precomputed distance table for speed — this is just num_subvectors table lookups + additions.

pub fn asymmetric_distance(&self, query: &[f32], codes: &[u8]) -> f32

Compute asymmetric distance between a query and an encoded vector.

pub fn search( &self, query: &[f32], all_codes: &[u8], tombstones: &[u8], k: usize, ) -> Vec<(usize, f32)>

Search a collection of PQ-encoded vectors and return the top-k nearest by asymmetric distance (smallest distance = most similar).

all_codes is a flat buffer: [n_vectors * num_subvectors]. tombstones marks deleted vectors.

pub fn search_rerank( &self, query: &[f32], all_codes: &[u8], vectors: &[Vec<f32>], tombstones: &[u8], candidates: usize, k: usize, ) -> Vec<(usize, f32)>

Search with PQ then re-rank top candidates with exact cosine similarity.

Returns (index, cosine_score) pairs sorted by score descending.

pub fn encode_all(&self, vectors: &[Vec<f32>]) -> Vec<u8>

Encode all vectors and return flat code buffer.

pub fn to_hdf5_data(&self) -> (&[f32], [i64; 3])

Serialize the quantizer state to flat data for HDF5 storage. Returns (codebook_flat, metadata: [num_subvectors, num_centroids, dim]).

pub fn from_hdf5_data(codebook: Vec<f32>, metadata: [i64; 3]) -> Self

Reconstruct from HDF5 data.