Skip to content

adriacabeza/adriANN

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 

History

37 Commits
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

adriANN

License

adriANN is an Approximate Nearest Neighbors library in Rust, based on SPANN: Highly-efficient Billion scale Approximate Nearest Neighbor Search. It aims to be:

  • Memory-Efficient Design: SPANN stores only the centroid points of posting lists in memory, significantly reducing memory requirements. This is important as most of the algorithms mainly focus on how to do low latency and high recall search all in memory with offline pre-built indexes. When targeting to the super large scale vector search scenarios, such as web search, the memory cost becomes extremely expensive.
  • Optimized Disk Access: Large posting lists are stored on disk. Each point is replicated into its L nearest centroids so boundary queries still find their neighbours without inflating posting lists.
  • High Recall, Low Latency: The in-memory head index is a kiddo KD-tree over centroids. The SPANN paper selects centroids via a BK-tree-based hierarchical balanced clustering procedure; we use plain random sampling instead. Empirically this produces indistinguishable recall on real-world datasets: large inputs are covered by the law of large numbers, and small inputs are insensitive to minor head-index imbalance.

Usage Examples

Building a Spann index:

use ndarray::{Array1, Array2};
use adriann::spann::config::Config;
use adriann::spann::spann_builder::SpannIndexBuilder;

fn main() {
    // 1. Load configuration from file
    // example_config.yaml
    //
    // clustering_params:
    //   distance_metric: Euclidean
    //   num_centroids: 4
    //   replica_count: 2
    // search_params:
    //   gamma: 0.2
    // output_path: data

    let config: Config =
        Config::from_file("examples/example_config.yaml").expect("Failed to load configuration");

    // 2. Prepare your dataset
    let data = Array2::from_shape_vec(
        (6, 2), 
        vec![
            1.0, 2.0,
            1.5, 2.5,
            8.0, 8.0,
            8.5, 8.5,
            4.0, 4.0,
            4.5, 4.5,
        ]
    ).unwrap();

    // 3. Build the index using a chosen dimension (here: 2).
    //    Omit .persist(...) in the builder (or leave output_path null in the
    //    YAML) for an in-memory-only build. Use .persist("path") or set
    //    output_path to use a mmap-backed on-disk index.
    let spann_index = SpannIndexBuilder::new(config)
        .with_data(data.view())
        .build::<2>()
        .expect("Failed to build SPANN index");

    println!("SPANN index built successfully!");
}

Querying the Spann index:

use ndarray::Array1;
use adriann::spann::config::Config;
use adriann::spann::spann_builder::SpannIndexBuilder;

fn main() {
    // 1. Load configuration
    let config: Config =
        Config::from_file("examples/example_config.yaml").expect("Failed to load configuration");

    // 2. Load the index from previously built data
    let spann_index = SpannIndexBuilder::new(config)
        .load::<2>()
        .expect("Failed to load SPANN index");

    // 3. Query your vector (query takes &[f32])
    let k = 1; // number of nearest neighbors
    let query_vector = Array1::from(vec![1.0_f32, 2.0]);
    let result = spann_index
        .query(query_vector.as_slice().unwrap(), k)
        .expect("search failed");
    // result: Vec<PointData>, sorted by distance ascending.
    // Each PointData carries { point_id, vector, distance }.

    // 4. Print result
    println!("Nearest neighbour: point_id: {:?} and vector: {:?}",
             result[0].point_id,
             result[0].vector);
}

Zero-alloc querying

For tight batch loops, query_into reuses a caller-provided buffer as the backing storage for its bounded max-heap β€” no per-call allocations other than the Vec<f32> clone of the final top-k vectors (and those reuse slot allocations across calls):

use adriann::spann::posting_lists::PointData;

let k = 10;
let mut buf: Vec<Option<PointData>> = (0..k).map(|_| None).collect();

for query in queries.iter() {
    let n = spann_index
        .query_into(query, &mut buf)
        .expect("search failed");
    for hit in buf[..n].iter().map(|o| o.as_ref().unwrap()) {
        // hit.point_id, hit.vector, hit.distance
    }
}

Performance

Release builds use.cargo/config.toml which passes -C target-cpu=native, so inner-loop SIMD (wide::f32x8 in src/distances/distance.rs) compiles down to AVX2/AVX-512/NEON as available on the host.

Numbers from cargo run --release --example bench_gist1m against ANN_GIST1M (1M Γ— 960-dim f32 vectors, 1000 queries, 100-NN ground truth) on the author's laptop. Clustering parameters: K=1000, L=8, num_centroids=Absolute(1000), in-memory index.

Build: ~35s clustering + ~10s posting-list materialization.

Query (k=10, query_into with a reused buffer, sweeping gamma):

gamma recall@10 p50 latency p95 latency throughput
0.05 0.80 5.5 ms 29.6 ms 107 qps
0.10 0.93 11.6 ms 48.4 ms 60 qps
0.20 0.99 29.6 ms 79.7 ms 26 qps
0.40 1.00 34.0 ms 76.6 ms 26 qps

query (which allocates per-call) is typically 15–40% slower than query_into at the same recall. See examples/bench_gist1m.rs for the full runner.

Configuration

A typical config.yaml might look like this:

clustering_params:
  # Supported distance metrics: Euclidean, Manhattan
  distance_metric: Euclidean
  # Pick one (or omit both for K = ceil(sqrt(n_points))):
  num_centroids: 1024
  # centroid_ratio: 0.05
  # How many centroids each point is replicated into (SPANN multi-assignment).
  replica_count: 8
  # Optional RNG seed for reproducible clustering.
  rng_seed: null

search_params:
  # (1 + gamma) * dist(q, c1) threshold for SPANN centroid gating.
  gamma: 0.2

output_path: data
  • distance_metric: How vectors are compared. Supported: Euclidean, Manhattan.
  • num_centroids / centroid_ratio: Number of head-index centroids, as an absolute count or a fraction of the dataset. Omitting both defaults to ceil(sqrt(n_points)).
  • replica_count: Each point is assigned to its L nearest centroids (SPANN's multi-assignment) for boundary coverage.
  • rng_seed: Optional seed for deterministic clustering.
  • search_params.gamma: The SPANN centroid-gating parameter at query time.
  • output_path (optional): Where the mmap-backed index is stored. If omitted, the index lives entirely in memory. The builder's .persist(path) / .in_memory() methods override this.

How does it work?

- Index Structure

The dataset is divided into groups called "posting lists" based on clusters of data points. Each group has a centroid, a kind of representative summary point for that group. These centroids are stored inside an in-memory index for fast access, while the actual data points in the groups are stored on disk.

- Search Process

When a query (like a question or search request) comes in, SPANN looks at the centroids in memory to quickly find the groups that are most relevant to the query. It then loads only those relevant groups from the disk into memory for a more detailed search to find the best matches.

- Optimizations

SPANN keeps the size of each group (posting list) manageable so it can load them quickly from disk. It also improves the groups by including additional nearby points to cover "boundary" cases where the best matches might be split across different groups. During the search, SPANN dynamically decides how many groups to check, ensuring a good balance between speed and accuracy.

References


Made with ❀ and posting lists.
We welcome contributions: issues, PRs, discussions.

About

An Approximate Nearest Neighbors library in Rust, based on SPANN, a highly-efficient billion scale aproximate nearest neighbor search πŸ’₯πŸ¦€

Topics

Resources

License

Stars

10 stars

Watchers

1 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages