flow_knn/
select.rs

1//! Structured KNN throughput records for method selection.
2//!
3//! Append-only JSONL lives at `data/knn_perf_matrix.jsonl`. Regenerate or extend
4//! with `cargo run -p flow-knn --example collect_matrix --features "hnsw,ann-search"`.
5
6use crate::config::{HnswParams, KnnMethod};
7use serde::{Deserialize, Serialize};
8
9/// One timed cell in the (n, d, method) performance matrix.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct PerfRecord {
12    /// Backend id: `exact`, `hnsw_usearch`, `hnsw_ann_search`, `exact_gpu`, `ivf_gpu`, `nndescent_gpu`.
13    pub method: String,
14    pub n: usize,
15    pub d: usize,
16    pub k: usize,
17    /// Median wall time in seconds for build+self-query (or exact pass).
18    pub median_secs: f64,
19    /// Elements/sec = n / median_secs.
20    pub throughput_elem_per_s: f64,
21    /// Optional host / CI label.
22    #[serde(default, skip_serializing_if = "Option::is_none")]
23    pub machine: Option<String>,
24    /// ISO-8601 capture time.
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub captured_at: Option<String>,
27}
28
29/// Preferences that bias automatic method choice.
30#[derive(Debug, Clone, Default)]
31pub struct RecommendOpts {
32    /// Prefer exact when it is within this factor of the best ANN (default 1.25).
33    pub exact_ok_factor: Option<f64>,
34    /// Force usearch even when ann-search is available (quantization / no faer).
35    pub prefer_usearch: bool,
36    /// Neighbours requested (affects exact vs ANN crossover slightly).
37    pub k: Option<usize>,
38    /// Include GPU method ids when the `gpu` feature is on and an adapter is up.
39    /// Default `false` so headless CI stays on CPU.
40    pub allow_gpu: bool,
41}
42
43/// Committed matrix shipped with the crate (`data/knn_perf_matrix.jsonl`).
44const SHIPPED_MATRIX_JSONL: &str = include_str!("../data/knn_perf_matrix.jsonl");
45
46/// Built-in snapshot used when JSONL parse fails.
47pub fn builtin_matrix() -> Vec<PerfRecord> {
48    parse_matrix_jsonl(SHIPPED_MATRIX_JSONL).unwrap_or_else(|_| {
49        // Hard-coded fallback from Criterion pacmap_knn grid (2026-07-23).
50        vec![
51            rec("hnsw_usearch", 50_000, 10, 60, 0.803),
52            rec("hnsw_ann_search", 50_000, 10, 60, 0.529),
53            rec("exact", 50_000, 10, 60, 1.265),
54            rec("hnsw_usearch", 100_000, 15, 60, 1.817),
55            rec("hnsw_ann_search", 100_000, 15, 60, 1.626),
56            rec("hnsw_usearch", 250_000, 20, 60, 9.254),
57            rec("hnsw_ann_search", 250_000, 20, 60, 5.757),
58            rec("hnsw_usearch", 500_000, 20, 60, 20.967),
59            rec("hnsw_ann_search", 500_000, 20, 60, 13.467),
60        ]
61    })
62}
63
64fn rec(method: &str, n: usize, d: usize, k: usize, median_secs: f64) -> PerfRecord {
65    PerfRecord {
66        method: method.to_string(),
67        n,
68        d,
69        k,
70        median_secs,
71        throughput_elem_per_s: n as f64 / median_secs,
72        machine: Some("local-criterion".into()),
73        captured_at: Some("2026-07-23".into()),
74    }
75}
76
77/// Parse JSONL text into records (skips blank / comment lines).
78pub fn parse_matrix_jsonl(text: &str) -> Result<Vec<PerfRecord>, String> {
79    let mut out = Vec::new();
80    for (i, line) in text.lines().enumerate() {
81        let line = line.trim();
82        if line.is_empty() || line.starts_with('#') {
83            continue;
84        }
85        let rec: PerfRecord = serde_json::from_str(line)
86            .map_err(|e| format!("jsonl line {}: {e}", i + 1))?;
87        out.push(rec);
88    }
89    Ok(out)
90}
91
92/// Load matrix: prefer `path` if readable, else the shipped JSONL / [`builtin_matrix`].
93pub fn load_matrix(path: Option<&std::path::Path>) -> Vec<PerfRecord> {
94    if let Some(p) = path
95        && let Ok(text) = std::fs::read_to_string(p)
96        && let Ok(recs) = parse_matrix_jsonl(&text)
97        && !recs.is_empty()
98    {
99        return recs;
100    }
101    builtin_matrix()
102}
103
104/// Recommend a [`KnnMethod`] for `n` points of dimension `d`.
105///
106/// Uses nearest measured cells in the performance matrix (log-distance in n×d),
107/// falling back to feature-aware heuristics when the matrix is sparse.
108pub fn recommend_method(n: usize, d: usize, opts: &RecommendOpts) -> KnnMethod {
109    recommend_method_with_matrix(n, d, opts, &load_matrix(None))
110}
111
112/// Same as [`recommend_method`] but with an explicit matrix (tests / custom datasets).
113pub fn recommend_method_with_matrix(
114    n: usize,
115    d: usize,
116    opts: &RecommendOpts,
117    matrix: &[PerfRecord],
118) -> KnnMethod {
119    let factor = opts.exact_ok_factor.unwrap_or(1.25);
120
121    // Very small problems: exact is fine and avoids index build.
122    if n <= 5_000 {
123        return KnnMethod::Exact;
124    }
125
126    let candidates = available_method_ids(opts.prefer_usearch, opts.allow_gpu, n);
127    let scored: Vec<(&str, f64)> = candidates
128        .iter()
129        .filter_map(|&method| estimate_secs(method, n, d, matrix).map(|secs| (method, secs)))
130        .collect();
131
132    if scored.is_empty() {
133        return heuristic_fallback(n, opts.prefer_usearch);
134    }
135
136    let best = scored
137        .iter()
138        .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
139        .copied();
140
141    if let Some((best_id, best_secs)) = best {
142        if best_id != "exact"
143            && let Some((_, exact_secs)) = scored.iter().find(|(id, _)| *id == "exact")
144            && *exact_secs <= best_secs * factor
145        {
146            return KnnMethod::Exact;
147        }
148        return method_from_id(best_id);
149    }
150
151    heuristic_fallback(n, opts.prefer_usearch)
152}
153
154fn available_method_ids(prefer_usearch: bool, allow_gpu: bool, n: usize) -> Vec<&'static str> {
155    let mut ids = Vec::new();
156    // Exact only competes up to FCS mid-scale; beyond that ANN build dominates.
157    if n <= 80_000 {
158        ids.push("exact");
159    }
160    #[cfg(feature = "hnsw")]
161    ids.push("hnsw_usearch");
162    #[cfg(feature = "ann-search")]
163    if !prefer_usearch {
164        ids.push("hnsw_ann_search");
165    }
166    #[cfg(feature = "gpu")]
167    if allow_gpu && crate::gpu_adapter_available() {
168        if n <= 50_000 {
169            ids.push("exact_gpu");
170        }
171        ids.push("ivf_gpu");
172        if n >= 50_000 {
173            ids.push("nndescent_gpu");
174        }
175    }
176    #[cfg(not(feature = "gpu"))]
177    let _ = allow_gpu;
178    ids
179}
180
181fn method_from_id(id: &str) -> KnnMethod {
182    match id {
183        "exact" => KnnMethod::Exact,
184        #[cfg(feature = "hnsw")]
185        "hnsw_usearch" => KnnMethod::Hnsw(HnswParams::default()),
186        #[cfg(feature = "ann-search")]
187        "hnsw_ann_search" => KnnMethod::AnnSearchHnsw(HnswParams::default()),
188        #[cfg(feature = "gpu")]
189        "exact_gpu" => KnnMethod::GpuExact,
190        #[cfg(feature = "gpu")]
191        "ivf_gpu" => KnnMethod::GpuIvf(crate::IvfGpuParams::default()),
192        #[cfg(feature = "gpu")]
193        "nndescent_gpu" => KnnMethod::GpuNnDescent(crate::NnDescentGpuParams::default()),
194        _ => heuristic_fallback(100_000, false),
195    }
196}
197
198fn heuristic_fallback(n: usize, prefer_usearch: bool) -> KnnMethod {
199    if n <= 5_000 {
200        return KnnMethod::Exact;
201    }
202    #[cfg(feature = "ann-search")]
203    if !prefer_usearch {
204        return KnnMethod::AnnSearchHnsw(HnswParams::default());
205    }
206    #[cfg(feature = "hnsw")]
207    {
208        return KnnMethod::Hnsw(HnswParams::default());
209    }
210    #[allow(unreachable_code)]
211    KnnMethod::Exact
212}
213
214/// Inverse-distance weighted estimate of median seconds at (n,d) for `method`.
215fn estimate_secs(method: &str, n: usize, d: usize, matrix: &[PerfRecord]) -> Option<f64> {
216    let mut num = 0.0;
217    let mut den = 0.0;
218    for r in matrix.iter().filter(|r| r.method == method) {
219        let dn = (n as f64).ln() - (r.n as f64).ln();
220        let dd = (d as f64).ln() - (r.d as f64).ln();
221        let dist = (dn * dn + dd * dd).sqrt();
222        let w = if dist < 1e-9 {
223            return Some(r.median_secs * (n as f64 / r.n as f64));
224        } else {
225            1.0 / (dist * dist)
226        };
227        // Scale observed time roughly with n (and mildly with d).
228        let scaled = r.median_secs * (n as f64 / r.n as f64) * ((d as f64 / r.d as f64).sqrt());
229        num += w * scaled;
230        den += w;
231    }
232    if den > 0.0 {
233        Some(num / den)
234    } else {
235        None
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242
243    #[test]
244    fn small_n_picks_exact() {
245        let m = recommend_method(1_000, 10, &RecommendOpts::default());
246        assert!(matches!(m, KnnMethod::Exact));
247    }
248
249    #[test]
250    fn large_n_prefers_ann_when_available() {
251        let m = recommend_method(250_000, 20, &RecommendOpts::default());
252        #[cfg(feature = "ann-search")]
253        assert!(matches!(m, KnnMethod::AnnSearchHnsw(_)));
254        #[cfg(all(not(feature = "ann-search"), feature = "hnsw"))]
255        assert!(matches!(m, KnnMethod::Hnsw(_)));
256    }
257
258    #[test]
259    fn prefer_usearch_honored() {
260        let opts = RecommendOpts {
261            prefer_usearch: true,
262            ..Default::default()
263        };
264        let m = recommend_method(250_000, 20, &opts);
265        #[cfg(feature = "hnsw")]
266        assert!(matches!(m, KnnMethod::Hnsw(_)));
267    }
268
269    #[test]
270    fn shipped_matrix_parses() {
271        let recs = builtin_matrix();
272        assert!(recs.len() >= 5);
273        assert!(recs.iter().any(|r| r.method == "hnsw_ann_search"));
274    }
275
276    #[test]
277    fn jsonl_roundtrip() {
278        let text = concat!(
279            r#"{"method":"exact","n":1000,"d":10,"k":10,"median_secs":0.01,"throughput_elem_per_s":100000}"#,
280            "\n"
281        );
282        let recs = parse_matrix_jsonl(text).unwrap();
283        assert_eq!(recs.len(), 1);
284        assert_eq!(recs[0].method, "exact");
285    }
286}
flow_knn/select.rs

flow_knn/
select.rs
