Skip to main content

flow_knn/
config.rs

1//! Configuration for k-NN search.
2
3/// Approximate / exact nearest-neighbour method.
4#[derive(Debug, Clone)]
5pub enum KnnMethod {
6    /// HNSW via usearch (C++ FFI, hardware SIMD) — default when the `hnsw` feature is on.
7    #[cfg(feature = "hnsw")]
8    Hnsw(HnswParams),
9
10    /// Exact brute-force O(n · k · d) with a bounded heap per query.
11    Exact,
12
13    /// k-d tree via `kiddo` (currently falls back to exact).
14    #[cfg(feature = "kdtree")]
15    KdTree,
16
17    /// HNSW via `ann-search-rs` (optional feature `ann-search`).
18    #[cfg(feature = "ann-search")]
19    AnnSearchHnsw(HnswParams),
20
21    /// Exact exhaustive kNN on GPU (ann-search-rs + cubeCL / wgpu). Feature `gpu`.
22    #[cfg(feature = "gpu")]
23    GpuExact,
24
25    /// IVF approximate kNN on GPU. Feature `gpu`.
26    #[cfg(feature = "gpu")]
27    GpuIvf(IvfGpuParams),
28
29    /// NN-Descent / CAGRA-style approximate kNN on GPU. Feature `gpu`.
30    #[cfg(feature = "gpu")]
31    GpuNnDescent(NnDescentGpuParams),
32
33    /// Reserved placeholder.
34    Annoy,
35}
36
37/// IVF-GPU list / probe parameters (`None` → ann-search-rs defaults √n / √nlist).
38#[derive(Debug, Clone, Default)]
39#[cfg(feature = "gpu")]
40pub struct IvfGpuParams {
41    pub n_list: Option<usize>,
42    pub n_probes: Option<usize>,
43}
44
45/// NN-Descent GPU graph parameters (`None` → library defaults).
46#[derive(Debug, Clone)]
47#[cfg(feature = "gpu")]
48pub struct NnDescentGpuParams {
49    /// Final graph degree after pruning (default: max(k, 30) at call site).
50    pub k: Option<usize>,
51    pub k_build: Option<usize>,
52    pub n_trees: Option<usize>,
53    pub delta: f32,
54    pub rho: Option<f32>,
55}
56
57#[cfg(feature = "gpu")]
58impl Default for NnDescentGpuParams {
59    fn default() -> Self {
60        Self {
61            k: None,
62            k_build: None,
63            n_trees: None,
64            delta: 0.001,
65            rho: None,
66        }
67    }
68}
69
70impl Default for KnnMethod {
71    fn default() -> Self {
72        // Prefer ann-search-rs when that feature is enabled (matches manifolds-rs HNSW).
73        #[cfg(feature = "ann-search")]
74        return Self::AnnSearchHnsw(HnswParams::default());
75        #[cfg(all(not(feature = "ann-search"), feature = "hnsw"))]
76        return Self::Hnsw(HnswParams::default());
77        #[cfg(all(not(feature = "ann-search"), not(feature = "hnsw")))]
78        return Self::Exact;
79    }
80}
81
82/// Quality / memory trade-off for HNSW indices (usearch or ann-search-rs).
83#[derive(Debug, Clone)]
84pub struct HnswParams {
85    /// Graph connectivity (M). Default 16.
86    pub m: usize,
87    /// Build-time candidate set size. Default 200.
88    pub ef_construction: usize,
89    /// Query-time candidate set size. Default 50.
90    pub ef_search: usize,
91    /// Storage quantization (usearch only; ignored by ann-search-rs).
92    pub quantization: Quantization,
93}
94
95impl Default for HnswParams {
96    fn default() -> Self {
97        Self {
98            m: 16,
99            ef_construction: 200,
100            ef_search: 50,
101            quantization: Quantization::F32,
102        }
103    }
104}
105
106/// Vector quantization for usearch HNSW storage.
107#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
108pub enum Quantization {
109    #[default]
110    F32,
111    F16,
112    I8,
113}
114
115/// Distance metric for the KNN graph.
116#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
117pub enum DistanceMetric {
118    #[default]
119    Euclidean,
120    EuclideanSq,
121    Cosine,
122    Manhattan,
123}