flat-knn 0.1.0

Simple and efficient SIMD-accelerated library for k-nearest neighbors search.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

# Flat KNN

Simple and efficient library for k-nearest neighbors search in Rust.

Library takes advantages of the SIMD capabilities of modern CPUs to perform fast k-NN search. It features a generic API that supports multiple data layouts, numeric types (`f32`, `f16`), and distance metrics without any runtime overhead.

## Installation

Add the following to your `Cargo.toml`. Currently, the library can be installed directly from the GitHub repository:

```toml
[dependencies]
flat-knn = { git = "https://github.com/simonplhak/flat-knn" }
```

## Usage

```rust
use flat_knn::{knn, Dot, L2};
use half::f16;

fn main() {
    let data = vec![
        1.0, 2.0, 3.0, 4.0, // 0
        5.0, 6.0, 7.0, 8.0, // 1
        9.0, 1.0, 2.0, 3.0, // 2
    ];
    let dim = 4;

    let query = [1.0, 2.0, 3.0, 5.0];
    let k = 2;
    
    // Compute L2
    let neighbors = knn::<_, L2>((&data, dim), &query, k);

    println!("Query: {query:?}");
    println!("{k} nearest neighbors (L2):");
    for (dist, index) in neighbors {
        println!("- Index: {index}, Distance: {dist}");
    }

    // Compute Inner Product
    let neighbors2 = knn::<_, Dot>((&data, dim), &query, k);

    println!("\n{k} nearest neighbors (Inner Product):");
    for (dist, index) in neighbors2 {
        println!("- Index: {index}, Distance: {dist}");
    }
    
    // Work with half-precision (f16) natively using the exact same API
    let data_f16: Vec<f16> = data.iter().map(|&x| f16::from_f32(x)).collect();
    let query_f16: Vec<f16> = query.iter().map(|&x| f16::from_f32(x)).collect();
    let neighbors_f16 = knn::<_, L2>((&data_f16, dim), &query_f16, k);

    println!("\n{k} nearest neighbors (L2 with f16):");
    for (dist, index) in neighbors_f16 {
        println!("- Index: {index}, Distance: {dist}");
    }
}
```

## Extensibility

The `flat-knn` crate relies strictly on static trait bounds, avoiding slow dynamic dispatch or abstraction overhead.

### Extending Distance Metrics

You can swap `L2` or `Dot` with your own distance calculations by implementing the `DistanceMetric<T>` trait on a custom struct.
This trait requires you to define how a calculation builds standard `(distance, index)` metadata into a sorting item for the generic max-heap/min-heap structure:

```rust
use flat_knn::DistanceMetric;
use ordered_float::OrderedFloat;

pub struct CustomMetric;

impl<T: VectorType> DistanceMetric<T> for CustomMetric {
    // Determine heap ordering. For example, smaller distances are better for L2:
    type HeapItem = (OrderedFloat<f32>, usize);

    #[inline(always)]
    fn build_item(query: &[T], chunk: &[T], i: usize) -> Self::HeapItem {
        let dist = /* do math here */;
        (OrderedFloat(dist), i)
    }

    #[inline(always)]
    fn extract_item(item: Self::HeapItem) -> (f32, usize) {
        (item.0.0, item.1)
    }
}
```

### Extending Data Types

By default, the crate implements the `Indexable<T>` trait for 1D vectors partitioned by stride `(Vec<T>, usize)`, nested `Vec<Vec<T>>`, and flat slices `&[T]`.
If you have a custom continuous memory structure like `ndarray` or matrices, you can easily provide a custom target struct implementing the `Indexable<T>` trait, defining `get`, `len`, and `is_empty`.