From f80afc085ffc3be3c186f9d46f041c99c41872f9 Mon Sep 17 00:00:00 2001 From: Giulio Rebuffo Date: Thu, 11 Dec 2025 16:20:41 +0100 Subject: [PATCH 1/8] save --- Cargo.toml | 1 + examples/lazy_memory_comparison.rs | 308 ++++++++++++++++++++++++++++ examples/recall.rs | 319 +++++++++++++++++++++-------- examples/recall_discrete.rs | 145 +++++++++++-- src/hnsw.rs | 2 + src/hnsw/feature_store.rs | 34 +++ src/hnsw/hnsw_const.rs | 98 ++++++--- src/lib.rs | 1 + tests/feature_store.rs | 277 +++++++++++++++++++++++++ 9 files changed, 1056 insertions(+), 129 deletions(-) create mode 100644 examples/lazy_memory_comparison.rs create mode 100644 src/hnsw/feature_store.rs create mode 100644 tests/feature_store.rs diff --git a/Cargo.toml b/Cargo.toml index e998192..c68d9f4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,7 @@ byteorder = "1.4.3" serde_json = "1.0.64" num-traits = "0.2.14" bitarray = { version = "0.9.1", default-features = false, features = ["space"] } +memmap2 = "0.9" [profile.dev] opt-level = 3 diff --git a/examples/lazy_memory_comparison.rs b/examples/lazy_memory_comparison.rs new file mode 100644 index 0000000..05214d3 --- /dev/null +++ b/examples/lazy_memory_comparison.rs @@ -0,0 +1,308 @@ +use hnsw::{FeatureStore, Hnsw, Params, Searcher}; +use rand::Rng; +use rand_pcg::Pcg64; +use space::{Metric, Neighbor}; +use std::convert::TryInto; +use std::fs::{File, OpenOptions}; +use std::io::{Read, Seek, SeekFrom, Write}; +use std::path::Path; +use std::time::Instant; + +struct Euclidean; + +impl Metric<[f32; 128]> for Euclidean { + type Unit = u32; + fn distance(&self, a: &[f32; 128], b: &[f32; 128]) -> u32 { + let sum: f32 = a + .iter() + .zip(b.iter()) + .map(|(&a, &b)| (a - b).powi(2)) + .sum(); + (sum.sqrt() * 1_000_000.0) as u32 + } +} + +const FEATURE_SIZE: usize = 128; +const FEATURE_BYTES: usize = FEATURE_SIZE * std::mem::size_of::(); + +/// A disk-based feature store that reads features from disk on every access. +/// This uses NO RAM for feature storage - only a small buffer for reading. +/// +/// Trade-off: Much slower due to disk I/O, but uses minimal memory. we are not +/// using mmapping here to isolate memory usage as low as possible (mmap is faster but +/// measurements are affected by OS page caching. realistically the results of this storage impl +/// show the true memory usage). +struct DiskFeatureStore { + file: File, + len: usize, +} + +impl DiskFeatureStore { + fn new>(path: P) -> std::io::Result { + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(path)?; + + Ok(Self { file, len: 0 }) + } +} + +impl FeatureStore<[f32; 128]> for DiskFeatureStore { + fn get(&self, index: usize) -> &[f32; 128] { + // Use thread-local storage to return a reference + // This is the key trick - we only keep ONE feature in memory at a time + thread_local! { + static BUFFER: std::cell::UnsafeCell<[f32; 128]> = std::cell::UnsafeCell::new([0.0; 128]); + } + + BUFFER.with(|buf| { + let buffer = unsafe { &mut *buf.get() }; + + let mut file = &self.file; + let offset = index * FEATURE_BYTES; + file.seek(SeekFrom::Start(offset as u64)).unwrap(); + + let mut bytes = [0u8; FEATURE_BYTES]; + file.read_exact(&mut bytes).unwrap(); + + for (i, chunk) in bytes.chunks_exact(4).enumerate() { + buffer[i] = f32::from_le_bytes(chunk.try_into().unwrap()); + } + + unsafe { &*(buffer as *const [f32; 128]) } + }) + } + + fn push(&mut self, feature: [f32; 128]) { + self.file.seek(SeekFrom::End(0)).unwrap(); + let bytes: Vec = feature.iter().flat_map(|f| f.to_le_bytes()).collect(); + self.file.write_all(&bytes).unwrap(); + self.len += 1; + } + + fn len(&self) -> usize { + self.len + } + + fn is_empty(&self) -> bool { + self.len == 0 + } +} + +/// A "null" feature store that doesn't store features at all. +/// Only useful for measuring the memory overhead of the graph structure alone. +struct NullFeatureStore { + len: usize, + dummy: [f32; 128], +} + +impl NullFeatureStore { + fn new() -> Self { + Self { + len: 0, + dummy: [0.0; 128], + } + } +} + +impl FeatureStore<[f32; 128]> for NullFeatureStore { + fn get(&self, _index: usize) -> &[f32; 128] { + // WARNING: This returns garbage! Only for memory measurement. + &self.dummy + } + + fn push(&mut self, _feature: [f32; 128]) { + self.len += 1; + } + + fn len(&self) -> usize { + self.len + } + + fn is_empty(&self) -> bool { + self.len == 0 + } +} + +fn generate_random_vectors(count: usize) -> Vec<[f32; 128]> { + let mut rng = rand::thread_rng(); + (0..count) + .map(|_| { + let mut v = [0.0f32; 128]; + for x in v.iter_mut() { + *x = rng.gen(); + } + v + }) + .collect() +} + +fn get_memory_usage() -> usize { + #[cfg(target_os = "linux")] + { + if let Ok(statm) = std::fs::read_to_string("/proc/self/statm") { + let parts: Vec<&str> = statm.split_whitespace().collect(); + if parts.len() >= 2 { + let rss_pages: usize = parts[1].parse().unwrap_or(0); + return rss_pages * 4096; + } + } + } + + #[cfg(target_os = "macos")] + { + use std::process::Command; + if let Ok(output) = Command::new("ps") + .args(&["-o", "rss=", "-p", &std::process::id().to_string()]) + .output() + { + if let Ok(rss_str) = String::from_utf8(output.stdout) { + if let Ok(rss_kb) = rss_str.trim().parse::() { + return rss_kb * 1024; + } + } + } + } + + 0 +} + +fn format_bytes(bytes: usize) -> String { + if bytes >= 1024 * 1024 * 1024 { + format!("{:.2} GB", bytes as f64 / (1024.0 * 1024.0 * 1024.0)) + } else if bytes >= 1024 * 1024 { + format!("{:.2} MB", bytes as f64 / (1024.0 * 1024.0)) + } else if bytes >= 1024 { + format!("{:.2} KB", bytes as f64 / 1024.0) + } else { + format!("{} bytes", bytes) + } +} + +fn main() -> std::io::Result<()> { + const NUM_VECTORS: usize = 100_000; + const NUM_QUERIES: usize = 10; + + let params = Params::new().ef_construction(20); // Lower for faster insertion + + // Generate test data + println!("Generating random vectors..."); + let vectors = generate_random_vectors(NUM_VECTORS); + let queries = generate_random_vectors(NUM_QUERIES); + + // Drop the generated vectors from consideration + let vectors_size = vectors.len() * FEATURE_BYTES; + + println!("=== Test 1: Graph structure only (NullFeatureStore) ==="); + + let mem_before_null = get_memory_usage(); + let graph_only_delta; + { + let storage = NullFeatureStore::new(); + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + + let mut hnsw: Hnsw = + Hnsw::new_with_storage_and_params(Euclidean, storage, params.clone(), prng); + let mut searcher = Searcher::default(); + + let start = Instant::now(); + for v in &vectors { + hnsw.insert(*v, &mut searcher); + } + + let mem_after = get_memory_usage(); + graph_only_delta = mem_after.saturating_sub(mem_before_null); + println!("Graph-only memory: +{}", format_bytes(graph_only_delta)); + } + println!(); + + println!("=== Test 2: In-memory Vec (default) ==="); + + let mem_before_vec = get_memory_usage(); + let vec_delta; + { + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + let mut hnsw: Hnsw = + Hnsw::new_params_and_prng(Euclidean, params.clone(), prng); + let mut searcher = Searcher::default(); + + let start = Instant::now(); + for v in &vectors { + hnsw.insert(*v, &mut searcher); + } + + let mem_after = get_memory_usage(); + vec_delta = mem_after.saturating_sub(mem_before_vec); + println!("Total memory: +{}", format_bytes(vec_delta)); + + let mut neighbors = [Neighbor { index: !0, distance: !0 }; 10]; + let start = Instant::now(); + for q in &queries { + hnsw.nearest(q, 64, &mut searcher, &mut neighbors); + } + println!("Search time: {:?} ({} queries)", start.elapsed(), NUM_QUERIES); + } + println!(); + + println!("=== Test 3: Disk-based FeatureStore ==="); + + let disk_path = "/tmp/hnsw_disk_features.bin"; + let mem_before_disk = get_memory_usage(); + let disk_delta; + { + let storage = DiskFeatureStore::new(disk_path)?; + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + + let mut hnsw: Hnsw = + Hnsw::new_with_storage_and_params(Euclidean, storage, params.clone(), prng); + let mut searcher = Searcher::default(); + + let start = Instant::now(); + for v in &vectors { + hnsw.insert(*v, &mut searcher); + } + + let mem_after = get_memory_usage(); + disk_delta = mem_after.saturating_sub(mem_before_disk); + + // Search (will be slow due to disk I/O) + let mut neighbors = [Neighbor { index: !0, distance: !0 }; 10]; + let start = Instant::now(); + for q in &queries { + hnsw.nearest(q, 64, &mut searcher, &mut neighbors); + } + println!("Search time: {:?} ({} queries) - slower due to disk I/O", start.elapsed(), NUM_QUERIES); + + std::fs::remove_file(disk_path).ok(); + } + println!(); + + println!("╔══════════════════════════════════════════════════════════════╗"); + println!("║ MEMORY SUMMARY ║"); + println!("╠══════════════════════════════════════════════════════════════╣"); + println!("║ Feature data size: {:>30} ║", format_bytes(NUM_VECTORS * FEATURE_BYTES)); + println!("╠══════════════════════════════════════════════════════════════╣"); + println!("║ Graph structure only: {:>30} ║", format_bytes(graph_only_delta)); + println!("║ Vec (graph + features): {:>30} ║", format_bytes(vec_delta)); + println!("║ Disk-based (graph only): {:>30} ║", format_bytes(disk_delta)); + println!("╠══════════════════════════════════════════════════════════════╣"); + + let feature_overhead = vec_delta.saturating_sub(graph_only_delta); + let savings = vec_delta.saturating_sub(disk_delta); + let savings_pct = if vec_delta > 0 { + (savings as f64 / vec_delta as f64) * 100.0 + } else { + 0.0 + }; + + println!("║ Feature storage overhead: {:>30} ║", format_bytes(feature_overhead)); + println!("║ Memory saved with disk: {:>30} ║", format_bytes(savings)); + println!("║ Savings percentage: {:>29.1}% ║", savings_pct); + println!("╚══════════════════════════════════════════════════════════════╝"); + println!(); + + Ok(()) +} diff --git a/examples/recall.rs b/examples/recall.rs index 99fcbd4..3e8b46b 100644 --- a/examples/recall.rs +++ b/examples/recall.rs @@ -1,12 +1,14 @@ use byteorder::{ByteOrder, LittleEndian}; use gnuplot::*; use hnsw::*; +use memmap2::MmapMut; use rand::distributions::Standard; use rand::{Rng, SeedableRng}; use rand_pcg::Pcg64; use space::Metric; use space::Neighbor; use std::cell::RefCell; +use std::fs::OpenOptions; use std::io::Read; use std::path::PathBuf; use structopt::StructOpt; @@ -25,37 +27,85 @@ impl Metric<&[f32]> for Euclidean { } } +impl Metric<[f32; 64]> for Euclidean { + type Unit = u32; + fn distance(&self, a: &[f32; 64], b: &[f32; 64]) -> u32 { + a.iter() + .zip(b.iter()) + .map(|(&a, &b)| (a - b).powi(2)) + .sum::() + .sqrt() + .to_bits() + } +} + +/// A mmap-based feature store for [f32; 64] arrays. +struct MmapFeatureStore { + mmap: MmapMut, + len: usize, + capacity: usize, +} + +impl MmapFeatureStore { + fn new(path: &str, capacity: usize) -> std::io::Result { + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(path)?; + + let total_size = capacity * std::mem::size_of::<[f32; 64]>(); + file.set_len(total_size as u64)?; + + let mmap = unsafe { MmapMut::map_mut(&file)? }; + + Ok(Self { + mmap, + len: 0, + capacity, + }) + } +} + +impl FeatureStore<[f32; 64]> for MmapFeatureStore { + fn get(&self, index: usize) -> &[f32; 64] { + let offset = index * std::mem::size_of::<[f32; 64]>(); + unsafe { &*(self.mmap[offset..].as_ptr() as *const [f32; 64]) } + } + + fn push(&mut self, feature: [f32; 64]) { + assert!(self.len < self.capacity, "MmapFeatureStore capacity exceeded"); + let offset = self.len * std::mem::size_of::<[f32; 64]>(); + let bytes: &[u8] = unsafe { + std::slice::from_raw_parts(feature.as_ptr() as *const u8, std::mem::size_of::<[f32; 64]>()) + }; + self.mmap[offset..offset + bytes.len()].copy_from_slice(bytes); + self.len += 1; + } + + fn len(&self) -> usize { + self.len + } + + fn is_empty(&self) -> bool { + self.len == 0 + } +} + #[derive(Debug, StructOpt)] #[structopt(name = "recall", about = "Generates recall graphs for HNSW")] struct Opt { /// The value of M to use. - /// - /// This can only be between 4 and 52 inclusive and a multiple of 4. - /// M0 is set to 2 * M. #[structopt(short = "m", long = "max_edges", default_value = "24")] m: usize, /// The dataset size to test on. #[structopt(short = "s", long = "size", default_value = "10000")] size: usize, /// Total number of query bitstrings. - /// - /// The higher this is, the better the quality of the output data and statistics, but - /// the longer the benchmark will take to set up. #[structopt(short = "q", long = "queries", default_value = "10000")] num_queries: usize, /// The number of dimensions in the feature vector. - /// - /// This is the length of the feature vector. The descriptor_stride (-d) - /// parameter must exceed this value. - /// - /// Possible values: - /// - 8 - /// - 16 - /// - 32 - /// - 64 - /// - 128 - /// - 256 - /// - 512 #[structopt(short = "l", long = "dimensions", default_value = "64")] dimensions: usize, /// The beginning ef value. @@ -71,14 +121,14 @@ struct Opt { #[structopt(short = "f", long = "file")] file: Option, /// The descriptor stride length in floats. - /// - /// KAZE: 64 - /// SIFT: 128 #[structopt(short = "d", long = "descriptor_stride", default_value = "64")] descriptor_stride: usize, /// efConstruction controlls the quality of the graph at build-time. #[structopt(short = "c", long = "ef_construction", default_value = "400")] ef_construction: usize, + /// Use mmap-based feature storage instead of in-memory Vec. + #[structopt(long = "mmap")] + mmap: bool, } fn process(opt: &Opt) -> (Vec, Vec) { @@ -91,54 +141,37 @@ fn process(opt: &Opt) -> (Vec, Vec) { let (search_space, query_strings): (Vec, Vec) = if let Some(filepath) = &opt.file { eprintln!( "Reading {} search space descriptors of size {} f32s from file \"{}\"...", - opt.size, - opt.descriptor_stride, - filepath.display() + opt.size, opt.descriptor_stride, filepath.display() ); let mut file = std::fs::File::open(filepath).expect("unable to open file"); - // We are loading floats, so multiply by 4. let mut search_space = vec![0u8; opt.size * opt.descriptor_stride * 4]; file.read_exact(&mut search_space).expect( "unable to read enough search descriptors from the file (try decreasing -s/-q)", ); - let search_space = search_space - .chunks_exact(4) - .map(LittleEndian::read_f32) - .collect(); + let search_space = search_space.chunks_exact(4).map(LittleEndian::read_f32).collect(); eprintln!("Done."); eprintln!( "Reading {} query descriptors of size {} f32s from file \"{}\"...", - opt.num_queries, - opt.descriptor_stride, - filepath.display() + opt.num_queries, opt.descriptor_stride, filepath.display() ); - // We are loading floats, so multiply by 4. let mut query_strings = vec![0u8; opt.num_queries * opt.descriptor_stride * 4]; file.read_exact(&mut query_strings) .expect("unable to read enough query descriptors from the file (try decreasing -q/-s)"); - let query_strings = query_strings - .chunks_exact(4) - .map(LittleEndian::read_f32) - .collect(); + let query_strings = query_strings.chunks_exact(4).map(LittleEndian::read_f32).collect(); eprintln!("Done."); (search_space, query_strings) } else { - eprintln!("Generating {} random bitstrings...", opt.size); + eprintln!("Generating {} random vectors...", opt.size); let search_space: Vec = rng .sample_iter(&Standard) .take(opt.size * opt.descriptor_stride) .collect(); eprintln!("Done."); - // Create another RNG to prevent potential correlation. let rng = Pcg64::from_seed([6; 32]); - - eprintln!( - "Generating {} independent random query strings...", - opt.num_queries - ); + eprintln!("Generating {} independent random query vectors...", opt.num_queries); let query_strings: Vec = rng .sample_iter(&Standard) .take(opt.num_queries * opt.descriptor_stride) @@ -156,10 +189,7 @@ fn process(opt: &Opt) -> (Vec, Vec) { .map(|c| &c[..opt.dimensions]) .collect(); - eprintln!( - "Computing the correct nearest neighbor distance for all {} queries...", - opt.num_queries - ); + eprintln!("Computing the correct nearest neighbor distance for all {} queries...", opt.num_queries); let correct_worst_distances: Vec<_> = query_strings .iter() .cloned() @@ -172,7 +202,6 @@ fn process(opt: &Opt) -> (Vec, Vec) { v.resize_with(opt.k, || unreachable!()); } } - // Get the worst distance v.into_iter().take(opt.k).last().unwrap() }) .collect(); @@ -195,21 +224,13 @@ fn process(opt: &Opt) -> (Vec, Vec) { let (recalls, times): (Vec, Vec) = efs .map(|ef| { let correct = RefCell::new(0usize); - let dest = vec![ - Neighbor { - index: !0, - distance: !0, - }; - opt.k - ]; + let dest = vec![Neighbor { index: !0, distance: !0 }; opt.k]; let stats = easybench::bench_env(dest, |mut dest| { let mut refmut = state.borrow_mut(); let (searcher, query) = &mut *refmut; let (ix, query_feature) = query.next().unwrap(); let correct_worst_distance = correct_worst_distances[ix]; - // Go through all the features. for &mut neighbor in hnsw.nearest(&query_feature, ef, searcher, &mut dest) { - // Any feature that is less than or equal to the worst real nearest neighbor distance is correct. if Euclidean.distance(&search_space[neighbor.index], &query_feature) <= correct_worst_distance { @@ -219,53 +240,175 @@ fn process(opt: &Opt) -> (Vec, Vec) { }); (stats, correct.into_inner()) }) - .fold( - (vec![], vec![]), - |(mut recalls, mut times), (stats, correct)| { - times.push((stats.ns_per_iter * 0.1f64.powi(9)).recip()); - // The maximum number of correct nearest neighbors is - recalls.push(correct as f64 / (stats.iterations * opt.k) as f64); - (recalls, times) - }, - ); + .fold((vec![], vec![]), |(mut recalls, mut times), (stats, correct)| { + times.push((stats.ns_per_iter * 0.1f64.powi(9)).recip()); + recalls.push(correct as f64 / (stats.iterations * opt.k) as f64); + (recalls, times) + }); eprintln!("Done."); (recalls, times) } +fn process_mmap, const M: usize, const M0: usize>( + opt: &Opt, + storage: S, +) -> (Vec, Vec) { + assert!(opt.k <= opt.size, "You must choose a dataset size larger or equal to the test search size"); + assert!(opt.dimensions == 64, "Mmap mode only supports 64 dimensions (use -l 64)"); + assert!(opt.file.is_none(), "Mmap mode does not support file input"); + + let rng = Pcg64::from_seed([5; 32]); + + eprintln!("Generating {} random vectors...", opt.size); + let raw: Vec = rng.clone().sample_iter(&Standard).take(opt.size * 64).collect(); + let search_space: Vec<[f32; 64]> = raw + .chunks_exact(64) + .map(|chunk| { + let mut arr = [0.0f32; 64]; + arr.copy_from_slice(chunk); + arr + }) + .collect(); + eprintln!("Done."); + + let rng = Pcg64::from_seed([6; 32]); + eprintln!("Generating {} independent random query vectors...", opt.num_queries); + let raw: Vec = rng.sample_iter(&Standard).take(opt.num_queries * 64).collect(); + let query_strings: Vec<[f32; 64]> = raw + .chunks_exact(64) + .map(|chunk| { + let mut arr = [0.0f32; 64]; + arr.copy_from_slice(chunk); + arr + }) + .collect(); + eprintln!("Done."); + + eprintln!("Computing the correct nearest neighbor distance for all {} queries...", opt.num_queries); + let correct_worst_distances: Vec<_> = query_strings + .iter() + .map(|feature| { + let mut v = vec![]; + for distance in search_space.iter().map(|n| Euclidean.distance(n, feature)) { + let pos = v.binary_search(&distance).unwrap_or_else(|e| e); + v.insert(pos, distance); + if v.len() > opt.k { + v.resize_with(opt.k, || unreachable!()); + } + } + v.into_iter().take(opt.k).last().unwrap() + }) + .collect(); + eprintln!("Done."); + + eprintln!("Generating HNSW with MmapFeatureStore..."); + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + let mut hnsw: Hnsw = + Hnsw::new_with_storage_and_params(Euclidean, storage, Params::new().ef_construction(opt.ef_construction), prng); + let mut searcher: Searcher<_> = Searcher::default(); + for feature in &search_space { + hnsw.insert(*feature, &mut searcher); + } + eprintln!("Done."); + + eprintln!("Computing recall graph..."); + let efs = opt.beginning_ef..=opt.ending_ef; + let state = RefCell::new((searcher, query_strings.iter().cloned().enumerate().cycle())); + let (recalls, times): (Vec, Vec) = efs + .map(|ef| { + let correct = RefCell::new(0usize); + let dest = vec![Neighbor { index: !0, distance: !0 }; opt.k]; + let stats = easybench::bench_env(dest, |mut dest| { + let mut refmut = state.borrow_mut(); + let (searcher, query) = &mut *refmut; + let (ix, query_feature) = query.next().unwrap(); + let correct_worst_distance = correct_worst_distances[ix]; + for &mut neighbor in hnsw.nearest(&query_feature, ef, searcher, &mut dest) { + if Euclidean.distance(&search_space[neighbor.index], &query_feature) + <= correct_worst_distance + { + *correct.borrow_mut() += 1; + } + } + }); + (stats, correct.into_inner()) + }) + .fold((vec![], vec![]), |(mut recalls, mut times), (stats, correct)| { + times.push((stats.ns_per_iter * 0.1f64.powi(9)).recip()); + recalls.push(correct as f64 / (stats.iterations * opt.k) as f64); + (recalls, times) + }); + eprintln!("Done."); + + (recalls, times) +} + +macro_rules! process_m { + ($opt:expr, $m:expr, $m0:expr) => { + process::<$m, $m0>(&$opt) + }; +} + +macro_rules! process_mmap_m { + ($opt:expr, $m:expr, $m0:expr) => { + process_mmap::<_, $m, $m0>(&$opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", $opt.size).unwrap()) + }; +} + fn main() { let opt = Opt::from_args(); - let (recalls, times) = { - // This can be increased indefinitely at the expense of compile time. - match opt.m { - 4 => process::<4, 8>(&opt), - 8 => process::<8, 16>(&opt), - 12 => process::<12, 24>(&opt), - 16 => process::<16, 32>(&opt), - 20 => process::<20, 40>(&opt), - 24 => process::<24, 48>(&opt), - 28 => process::<28, 56>(&opt), - 32 => process::<32, 64>(&opt), - 36 => process::<36, 72>(&opt), - 40 => process::<40, 80>(&opt), - 44 => process::<44, 88>(&opt), - 48 => process::<48, 96>(&opt), - 52 => process::<52, 104>(&opt), + let (recalls, times, storage_type) = if opt.mmap { + let (r, t) = match opt.m { + 4 => process_mmap_m!(opt, 4, 8), + 8 => process_mmap_m!(opt, 8, 16), + 12 => process_mmap_m!(opt, 12, 24), + 16 => process_mmap_m!(opt, 16, 32), + 20 => process_mmap_m!(opt, 20, 40), + 24 => process_mmap_m!(opt, 24, 48), + 28 => process_mmap_m!(opt, 28, 56), + 32 => process_mmap_m!(opt, 32, 64), + 36 => process_mmap_m!(opt, 36, 72), + 40 => process_mmap_m!(opt, 40, 80), + 44 => process_mmap_m!(opt, 44, 88), + 48 => process_mmap_m!(opt, 48, 96), + 52 => process_mmap_m!(opt, 52, 104), _ => { eprintln!("Only M between 4 and 52 inclusive and multiples of 4 are allowed"); return; } - } + }; + (r, t, "MmapFeatureStore") + } else { + let (r, t) = match opt.m { + 4 => process_m!(opt, 4, 8), + 8 => process_m!(opt, 8, 16), + 12 => process_m!(opt, 12, 24), + 16 => process_m!(opt, 16, 32), + 20 => process_m!(opt, 20, 40), + 24 => process_m!(opt, 24, 48), + 28 => process_m!(opt, 28, 56), + 32 => process_m!(opt, 32, 64), + 36 => process_m!(opt, 36, 72), + 40 => process_m!(opt, 40, 80), + 44 => process_m!(opt, 44, 88), + 48 => process_m!(opt, 48, 96), + 52 => process_m!(opt, 52, 104), + _ => { + eprintln!("Only M between 4 and 52 inclusive and multiples of 4 are allowed"); + return; + } + }; + (r, t, "Vec") }; let mut fg = Figure::new(); - fg.axes2d() .set_title( &format!( - "{}-NN Recall Graph (dimensions = {}, size = {}, M = {})", - opt.k, opt.dimensions, opt.size, opt.m + "{}-NN Recall Graph (dimensions = {}, size = {}, M = {}, storage = {})", + opt.k, opt.dimensions, opt.size, opt.m, storage_type ), &[], ) @@ -279,5 +422,5 @@ fn main() { .set_y_grid(true) .set_y_minor_grid(true); - fg.show().expect("unable to show gnuplot"); + fg.show().ok(); } diff --git a/examples/recall_discrete.rs b/examples/recall_discrete.rs index ff1d02d..9afb76e 100644 --- a/examples/recall_discrete.rs +++ b/examples/recall_discrete.rs @@ -2,16 +2,72 @@ use bitarray::{BitArray, Hamming}; use gnuplot::*; use hnsw::*; use itertools::Itertools; +use memmap2::MmapMut; use num_traits::Zero; use rand::distributions::Standard; use rand::{Rng, SeedableRng}; use rand_pcg::Pcg64; use space::*; use std::cell::RefCell; +use std::fs::OpenOptions; use std::io::Read; use std::path::PathBuf; use structopt::StructOpt; +/// A mmap-based feature store for BitArray types. +struct MmapBitArrayStore { + mmap: MmapMut, + len: usize, + capacity: usize, + _marker: std::marker::PhantomData, +} + +impl MmapBitArrayStore { + fn new(path: &str, capacity: usize, item_size: usize) -> std::io::Result { + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(path)?; + + let total_size = capacity * item_size; + file.set_len(total_size as u64)?; + + let mmap = unsafe { MmapMut::map_mut(&file)? }; + + Ok(Self { + mmap, + len: 0, + capacity, + _marker: std::marker::PhantomData, + }) + } +} + +impl FeatureStore> for MmapBitArrayStore> { + fn get(&self, index: usize) -> &BitArray { + let offset = index * N; + // Safety: BitArray is repr(transparent) over [u8; N] + unsafe { &*(self.mmap[offset..].as_ptr() as *const BitArray) } + } + + fn push(&mut self, feature: BitArray) { + assert!(self.len < self.capacity, "MmapBitArrayStore capacity exceeded"); + let offset = self.len * N; + self.mmap[offset..offset + N].copy_from_slice(feature.bytes()); + self.len += 1; + } + + fn len(&self) -> usize { + self.len + } + + fn is_empty(&self) -> bool { + self.len == 0 + } +} + #[derive(Debug, StructOpt)] #[structopt( name = "recall_discrete", @@ -69,11 +125,17 @@ struct Opt { /// efConstruction controlls the quality of the graph at build-time. #[structopt(short = "c", long = "ef_construction", default_value = "400")] ef_construction: usize, + /// Use disk-based feature storage instead of in-memory Vec. + /// This tests the FeatureStore trait with a disk backend. + /// Note: Does not support file input (-f). + #[structopt(long = "disk")] + disk: bool, } -fn process( +fn process, const M: usize, const M0: usize>( opt: &Opt, conv: fn(&[u8]) -> T, + storage: S, ) -> (Vec, Vec) where Hamming: Metric, @@ -164,8 +226,9 @@ where eprintln!("Done."); eprintln!("Generating HNSW..."); - let mut hnsw: Hnsw<_, T, Pcg64, M, M0> = - Hnsw::new_params(Hamming, Params::new().ef_construction(opt.ef_construction)); + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + let mut hnsw: Hnsw = + Hnsw::new_with_storage_and_params(Hamming, storage, Params::new().ef_construction(opt.ef_construction), prng); let mut searcher: Searcher<_> = Searcher::default(); for feature in &search_space { hnsw.insert(feature.clone(), &mut searcher); @@ -219,27 +282,56 @@ where macro_rules! process_m { ( $opt:expr, $m:expr, $m0:expr ) => { match $opt.bitstring_length { - 128 => process::, $m, $m0>(&$opt, |b| { + 128 => process::, _, $m, $m0>(&$opt, |b| { + let mut arr = [0; 16]; + for (d, &s) in arr.iter_mut().zip(b) { + *d = s; + } + BitArray::new(arr) + }, Vec::new()), + 256 => process::, _, $m, $m0>(&$opt, |b| { + let mut arr = [0; 32]; + for (d, &s) in arr.iter_mut().zip(b) { + *d = s; + } + BitArray::new(arr) + }, Vec::new()), + 512 => process::, _, $m, $m0>(&$opt, |b| { + let mut arr = [0; 64]; + for (d, &s) in arr.iter_mut().zip(b) { + *d = s; + } + BitArray::new(arr) + }, Vec::new()), + _ => panic!("error: incorrect bitstring_length, see --help for choices"), + } + }; +} + +macro_rules! process_mmap_m { + ( $opt:expr, $m:expr, $m0:expr ) => { + match $opt.bitstring_length { + 128 => process::, _, $m, $m0>(&$opt, |b| { let mut arr = [0; 16]; for (d, &s) in arr.iter_mut().zip(b) { *d = s; } BitArray::new(arr) - }), - 256 => process::, $m, $m0>(&$opt, |b| { + }, MmapBitArrayStore::new("/tmp/recall_discrete_mmap.bin", $opt.size, 16).unwrap()), + 256 => process::, _, $m, $m0>(&$opt, |b| { let mut arr = [0; 32]; for (d, &s) in arr.iter_mut().zip(b) { *d = s; } BitArray::new(arr) - }), - 512 => process::, $m, $m0>(&$opt, |b| { + }, MmapBitArrayStore::new("/tmp/recall_discrete_mmap.bin", $opt.size, 32).unwrap()), + 512 => process::, _, $m, $m0>(&$opt, |b| { let mut arr = [0; 64]; for (d, &s) in arr.iter_mut().zip(b) { *d = s; } BitArray::new(arr) - }), + }, MmapBitArrayStore::new("/tmp/recall_discrete_mmap.bin", $opt.size, 64).unwrap()), _ => panic!("error: incorrect bitstring_length, see --help for choices"), } }; @@ -248,9 +340,29 @@ macro_rules! process_m { fn main() { let opt = Opt::from_args(); - let (recalls, times) = { - // This can be increased indefinitely at the expense of compile time. - match opt.m { + let (recalls, times, storage_type) = if opt.disk { + let (r, t) = match opt.m { + 4 => process_mmap_m!(opt, 4, 8), + 8 => process_mmap_m!(opt, 8, 16), + 12 => process_mmap_m!(opt, 12, 24), + 16 => process_mmap_m!(opt, 16, 32), + 20 => process_mmap_m!(opt, 20, 40), + 24 => process_mmap_m!(opt, 24, 48), + 28 => process_mmap_m!(opt, 28, 56), + 32 => process_mmap_m!(opt, 32, 64), + 36 => process_mmap_m!(opt, 36, 72), + 40 => process_mmap_m!(opt, 40, 80), + 44 => process_mmap_m!(opt, 44, 88), + 48 => process_mmap_m!(opt, 48, 96), + 52 => process_mmap_m!(opt, 52, 104), + _ => { + eprintln!("Only M between 4 and 52 inclusive and multiples of 4 are allowed"); + return; + } + }; + (r, t, "MmapBitArrayStore") + } else { + let (r, t) = match opt.m { 4 => process_m!(opt, 4, 8), 8 => process_m!(opt, 8, 16), 12 => process_m!(opt, 12, 24), @@ -268,7 +380,8 @@ fn main() { eprintln!("Only M between 4 and 52 inclusive and multiples of 4 are allowed"); return; } - } + }; + (r, t, "Vec") }; let mut fg = Figure::new(); @@ -276,8 +389,8 @@ fn main() { fg.axes2d() .set_title( &format!( - "{}-NN Recall Graph (bits = {}, size = {}, M = {})", - opt.k, opt.bitstring_length, opt.size, opt.m + "{}-NN Recall Graph (bits = {}, size = {}, M = {}, storage = {})", + opt.k, opt.bitstring_length, opt.size, opt.m, storage_type ), &[], ) @@ -291,5 +404,5 @@ fn main() { .set_y_grid(true) .set_y_minor_grid(true); - fg.show().expect("unable to show gnuplot"); + fg.show().ok(); // Don't panic if gnuplot unavailable } diff --git a/src/hnsw.rs b/src/hnsw.rs index 032e817..f398295 100644 --- a/src/hnsw.rs +++ b/src/hnsw.rs @@ -1,6 +1,8 @@ +mod feature_store; mod hnsw_const; mod nodes; #[cfg(feature = "serde")] mod serde_impl; +pub use feature_store::FeatureStore; pub use hnsw_const::*; diff --git a/src/hnsw/feature_store.rs b/src/hnsw/feature_store.rs new file mode 100644 index 0000000..c018f20 --- /dev/null +++ b/src/hnsw/feature_store.rs @@ -0,0 +1,34 @@ +use alloc::vec::Vec; + +/// Trait for abstracting feature storage in HNSW. +pub trait FeatureStore { + fn get(&self, index: usize) -> &T; + fn push(&mut self, feature: T); + fn len(&self) -> usize; + #[inline] + fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +impl FeatureStore for Vec { + #[inline] + fn get(&self, index: usize) -> &T { + &self[index] + } + + #[inline] + fn push(&mut self, feature: T) { + Vec::push(self, feature) + } + + #[inline] + fn len(&self) -> usize { + Vec::len(self) + } + + #[inline] + fn is_empty(&self) -> bool { + Vec::is_empty(self) + } +} diff --git a/src/hnsw/hnsw_const.rs b/src/hnsw/hnsw_const.rs index 630b3b8..3fd27eb 100644 --- a/src/hnsw/hnsw_const.rs +++ b/src/hnsw/hnsw_const.rs @@ -1,7 +1,9 @@ +use super::feature_store::FeatureStore; use super::nodes::{HasNeighbors, Layer}; use crate::hnsw::nodes::{NeighborNodes, Node}; use crate::*; use alloc::{vec, vec::Vec}; +use core::marker::PhantomData; use num_traits::Zero; use rand_core::{RngCore, SeedableRng}; #[cfg(feature = "serde")] @@ -11,16 +13,19 @@ use space::{Knn, KnnPoints, Metric, Neighbor}; /// This provides a HNSW implementation for any distance function. /// /// The type `T` must implement [`space::Metric`] to get implementations. +/// +/// The type `S` is the feature storage backend, defaulting to `Vec`. +/// Custom storage backends can be used by implementing [`FeatureStore`]. #[derive(Clone)] #[cfg_attr( feature = "serde", derive(Serialize, Deserialize), serde(bound( - serialize = "Met: Serialize, T: Serialize, R: Serialize", - deserialize = "Met: Deserialize<'de>, T: Deserialize<'de>, R: Deserialize<'de>" + serialize = "Met: Serialize, T: Serialize, R: Serialize, S: Serialize", + deserialize = "Met: Deserialize<'de>, T: Deserialize<'de>, R: Deserialize<'de>, S: Deserialize<'de>" )) )] -pub struct Hnsw { +pub struct Hnsw = Vec> { /// Contains the space metric. metric: Met, /// Contains the zero layer. @@ -28,16 +33,18 @@ pub struct Hnsw { /// Contains the features of the zero layer. /// These are stored separately to allow SIMD speedup in the future by /// grouping small worlds of features together. - features: Vec, + features: S, /// Contains each non-zero layer. layers: Vec>>, /// This needs to create resonably random outputs to determine the levels of insertions. prng: R, /// The parameters for the HNSW. params: Params, + /// Marker for the feature type T. + _marker: PhantomData, } -impl Hnsw +impl Hnsw> where R: RngCore + SeedableRng, { @@ -50,6 +57,7 @@ where layers: vec![], prng: R::from_seed(R::Seed::default()), params: Params::new(), + _marker: PhantomData, } } @@ -62,6 +70,7 @@ where layers: vec![], prng: R::from_seed(R::Seed::default()), params, + _marker: PhantomData, } } @@ -73,14 +82,16 @@ where layers: vec![], prng: R::from_seed(R::Seed::default()), params, + _marker: PhantomData, } } } -impl Knn for Hnsw +impl Knn for Hnsw where R: RngCore, Met: Metric, + S: FeatureStore, { type Ix = usize; type Metric = Met; @@ -104,17 +115,18 @@ where } } -impl KnnPoints for Hnsw +impl KnnPoints for Hnsw where R: RngCore, Met: Metric, + S: FeatureStore, { fn get_point(&self, index: usize) -> &'_ T { - &self.features[index] + self.features.get(index) } } -impl Hnsw +impl Hnsw> where R: RngCore, Met: Metric, @@ -128,6 +140,7 @@ where layers: vec![], prng, params: Default::default(), + _marker: PhantomData, } } @@ -140,6 +153,40 @@ where layers: vec![], prng, params, + _marker: PhantomData, + } + } +} + +impl Hnsw +where + R: RngCore, + Met: Metric, + S: FeatureStore, +{ + /// Creates a HNSW with a custom feature storage backend. + pub fn new_with_storage(metric: Met, storage: S, prng: R) -> Self { + Self { + metric, + zero: vec![], + features: storage, + layers: vec![], + prng, + params: Default::default(), + _marker: PhantomData, + } + } + + /// Creates a HNSW with a custom feature storage backend and params. + pub fn new_with_storage_and_params(metric: Met, storage: S, params: Params, prng: R) -> Self { + Self { + metric, + zero: vec![], + features: storage, + layers: vec![], + prng, + params, + _marker: PhantomData, } } @@ -238,12 +285,12 @@ where /// /// The `item` must be retrieved from [`HNSW::search_layer`]. pub fn feature(&self, item: usize) -> &T { - &self.features[item as usize] + self.features.get(item) } /// Extract the feature from a particular level for a given item returned by [`HNSW::search_layer`]. pub fn layer_feature(&self, level: usize, item: usize) -> &T { - &self.features[self.layer_item_id(level, item) as usize] + self.features.get(self.layer_item_id(level, item)) } /// Retrieve the item ID for a given layer item returned by [`HNSW::search_layer`]. @@ -265,7 +312,7 @@ where pub fn layer_len(&self, level: usize) -> usize { if level == 0 { - self.features.len() + FeatureStore::len(&self.features) } else if level < self.layers() { self.layers[level - 1].len() } else { @@ -306,7 +353,7 @@ where self.search_single_layer(q, searcher, Layer::NonZero(layer), cap); if ix + 1 == level { let found = core::cmp::min(dest.len(), searcher.nearest.len()); - dest.copy_from_slice(&searcher.nearest[..found]); + dest[..found].copy_from_slice(&searcher.nearest[..found]); return &mut dest[..found]; } self.lower_search(layer, searcher); @@ -318,7 +365,7 @@ where self.search_zero_layer(q, searcher, cap); let found = core::cmp::min(dest.len(), searcher.nearest.len()); - dest.copy_from_slice(&searcher.nearest[..found]); + dest[..found].copy_from_slice(&searcher.nearest[..found]); &mut dest[..found] } @@ -348,7 +395,7 @@ where // Compute the distance of this neighbor. let distance = self .metric - .distance(q, &self.features[node_to_visit as usize]); + .distance(q, self.features.get(node_to_visit)); // Attempt to insert into nearest queue. let pos = searcher.nearest.partition_point(|n| n.distance <= distance); if pos != cap { @@ -423,9 +470,9 @@ where /// Gets the entry point's feature. fn entry_feature(&self) -> &T { if let Some(last_layer) = self.layers.last() { - &self.features[last_layer[0].zero_node as usize] + self.features.get(last_layer[0].zero_node) } else { - &self.features[0] + self.features.get(0) } } @@ -477,13 +524,13 @@ where // This is different for the zero layer. let (target_feature, target_neighbors) = if layer == 0 { ( - &self.features[target_ix], + self.features.get(target_ix), &self.zero[target_ix].neighbors[..], ) } else { let target = &self.layers[layer - 1][target_ix]; ( - &self.features[target.zero_node], + self.features.get(target.zero_node), &target.neighbors.neighbors[..], ) }; @@ -511,13 +558,14 @@ where None } else { // Compute the distance. The feature is looked up differently for the zero layer. + let neighbor_zero_node = if layer == 0 { + n + } else { + self.layers[layer - 1][n].zero_node + }; let distance = self.metric.distance( target_feature, - &self.features[if layer == 0 { - n - } else { - self.layers[layer - 1][n].zero_node - }], + self.features.get(neighbor_zero_node), ); Some((ix, distance)) } @@ -541,7 +589,7 @@ where } } -impl Default for Hnsw +impl Default for Hnsw> where R: RngCore + SeedableRng, Met: Default, diff --git a/src/lib.rs b/src/lib.rs index 33b9ff2..e105ca5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ extern crate alloc; mod hnsw; +pub use self::hnsw::FeatureStore; pub use self::hnsw::*; use ahash::RandomState; diff --git a/tests/feature_store.rs b/tests/feature_store.rs new file mode 100644 index 0000000..d97779f --- /dev/null +++ b/tests/feature_store.rs @@ -0,0 +1,277 @@ +//! Tests for the FeatureStore trait and custom storage backends. + +use hnsw::{FeatureStore, Hnsw, Searcher}; +use rand_pcg::Pcg64; +use space::{Metric, Neighbor}; +use std::cell::Cell; + +struct Euclidean; + +impl Metric<[f64; 4]> for Euclidean { + type Unit = u64; + fn distance(&self, a: &[f64; 4], b: &[f64; 4]) -> u64 { + a.iter() + .zip(b.iter()) + .map(|(&a, &b)| (a - b).powi(2)) + .sum::() + .sqrt() + .to_bits() + } +} + +struct TrackedFeatureStore { + inner: Vec, + get_count: Cell, + push_count: Cell, +} + +impl TrackedFeatureStore { + fn new() -> Self { + Self { + inner: Vec::new(), + get_count: Cell::new(0), + push_count: Cell::new(0), + } + } + + fn get_count(&self) -> usize { + self.get_count.get() + } + + fn push_count(&self) -> usize { + self.push_count.get() + } + + fn reset_counts(&self) { + self.get_count.set(0); + self.push_count.set(0); + } +} + +impl FeatureStore for TrackedFeatureStore { + fn get(&self, index: usize) -> &T { + self.get_count.set(self.get_count.get() + 1); + &self.inner[index] + } + + fn push(&mut self, feature: T) { + self.push_count.set(self.push_count.get() + 1); + self.inner.push(feature); + } + + fn len(&self) -> usize { + self.inner.len() + } + + fn is_empty(&self) -> bool { + self.inner.is_empty() + } +} + +// A simple in-memory cache that simulates lazy loading behavior. +struct MockLazyStore { + storage: Vec>, + loaded_count: Cell, +} + +impl MockLazyStore { + fn new() -> Self { + Self { + storage: Vec::new(), + loaded_count: Cell::new(0), + } + } + + fn loaded_count(&self) -> usize { + self.loaded_count.get() + } +} + +impl FeatureStore for MockLazyStore { + fn get(&self, index: usize) -> &T { + self.loaded_count.set(self.loaded_count.get() + 1); + self.storage[index].as_ref().expect("Feature not found") + } + + fn push(&mut self, feature: T) { + self.storage.push(Some(feature)); + } + + fn len(&self) -> usize { + self.storage.len() + } + + fn is_empty(&self) -> bool { + self.storage.is_empty() + } +} + +#[test] +fn test_tracked_feature_store() { + let storage = TrackedFeatureStore::<[f64; 4]>::new(); + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + + let mut hnsw: Hnsw> = + Hnsw::new_with_storage(Euclidean, storage, prng); + + let mut searcher = Searcher::default(); + + // Insert features + let features = [ + [0.0, 0.0, 0.0, 1.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + ]; + + for feature in features.iter() { + hnsw.insert(*feature, &mut searcher); + } + + // Verify push was called for each insert + assert_eq!(hnsw.len(), 4); + + // Search and verify get is called + let mut neighbors = [Neighbor { + index: !0, + distance: !0, + }; 4]; + + hnsw.nearest(&[0.0, 0.0, 0.0, 1.0], 24, &mut searcher, &mut neighbors); + + // First result should be exact match + assert_eq!(neighbors[0].index, 0); + + // Verify feature() method works through the trait + let feature = hnsw.feature(0); + assert_eq!(*feature, [0.0, 0.0, 0.0, 1.0]); +} + +#[test] +fn test_mock_lazy_store() { + // Test a mock lazy loading store + let storage = MockLazyStore::<[f64; 4]>::new(); + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + + let mut hnsw: Hnsw> = + Hnsw::new_with_storage(Euclidean, storage, prng); + + let mut searcher = Searcher::default(); + + // Insert features + let features = [ + [0.0, 0.0, 0.0, 1.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 1.0], + [0.0, 1.0, 1.0, 0.0], + [1.0, 1.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 1.0], + ]; + + for feature in features.iter() { + hnsw.insert(*feature, &mut searcher); + } + + assert_eq!(hnsw.len(), 8); + + // Search + let mut neighbors = [Neighbor { + index: !0, + distance: !0, + }; 4]; + + hnsw.nearest(&[0.0, 0.0, 0.0, 1.0], 24, &mut searcher, &mut neighbors); + + // Verify correct results + assert_eq!(neighbors[0].index, 0); // Exact match + assert_eq!(neighbors[0].distance, 0); +} + +#[test] +fn test_feature_store_with_params() { + // Test new_with_storage_and_params constructor + use hnsw::Params; + + let storage = TrackedFeatureStore::<[f64; 4]>::new(); + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + let params = Params::new().ef_construction(100); + + let mut hnsw: Hnsw> = + Hnsw::new_with_storage_and_params(Euclidean, storage, params, prng); + + let mut searcher = Searcher::default(); + + hnsw.insert([1.0, 2.0, 3.0, 4.0], &mut searcher); + assert_eq!(hnsw.len(), 1); +} + +#[test] +fn test_feature_store_layer_feature() { + // Test layer_feature method works with custom store + let storage = TrackedFeatureStore::<[f64; 4]>::new(); + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + + let mut hnsw: Hnsw> = + Hnsw::new_with_storage(Euclidean, storage, prng); + + let mut searcher = Searcher::default(); + + let features = [ + [0.0, 0.0, 0.0, 1.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + ]; + + for feature in features.iter() { + hnsw.insert(*feature, &mut searcher); + } + + // Test layer_feature at level 0 + let feature = hnsw.layer_feature(0, 0); + assert_eq!(*feature, [0.0, 0.0, 0.0, 1.0]); + + let feature = hnsw.layer_feature(0, 1); + assert_eq!(*feature, [0.0, 0.0, 1.0, 0.0]); +} + +#[test] +fn test_feature_store_empty_checks() { + // Test is_empty and len on custom store + let storage = TrackedFeatureStore::<[f64; 4]>::new(); + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + + let mut hnsw: Hnsw> = + Hnsw::new_with_storage(Euclidean, storage, prng); + + assert!(hnsw.is_empty()); + assert_eq!(hnsw.len(), 0); + assert_eq!(hnsw.layer_len(0), 0); + + let mut searcher = Searcher::default(); + hnsw.insert([1.0, 2.0, 3.0, 4.0], &mut searcher); + + assert!(!hnsw.is_empty()); + assert_eq!(hnsw.len(), 1); + assert_eq!(hnsw.layer_len(0), 1); +} + +#[test] +fn test_search_empty_custom_store() { + // Test searching an empty HNSW with custom store returns empty results + let storage = TrackedFeatureStore::<[f64; 4]>::new(); + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + + let hnsw: Hnsw> = + Hnsw::new_with_storage(Euclidean, storage, prng); + + let mut searcher = Searcher::default(); + let mut neighbors = [Neighbor { + index: !0, + distance: !0, + }; 4]; + + let result = hnsw.nearest(&[0.0, 0.0, 0.0, 1.0], 24, &mut searcher, &mut neighbors); + assert_eq!(result.len(), 0); +} From d07b9636cece16ed6d11ca99cb4af8a981f5e923 Mon Sep 17 00:00:00 2001 From: Giulio Rebuffo Date: Thu, 11 Dec 2025 19:04:03 +0100 Subject: [PATCH 2/8] save --- examples/lazy_memory_comparison.rs | 308 ----------------------------- examples/memory_usage.rs | 231 ++++++++++++++++++++++ 2 files changed, 231 insertions(+), 308 deletions(-) delete mode 100644 examples/lazy_memory_comparison.rs create mode 100644 examples/memory_usage.rs diff --git a/examples/lazy_memory_comparison.rs b/examples/lazy_memory_comparison.rs deleted file mode 100644 index 05214d3..0000000 --- a/examples/lazy_memory_comparison.rs +++ /dev/null @@ -1,308 +0,0 @@ -use hnsw::{FeatureStore, Hnsw, Params, Searcher}; -use rand::Rng; -use rand_pcg::Pcg64; -use space::{Metric, Neighbor}; -use std::convert::TryInto; -use std::fs::{File, OpenOptions}; -use std::io::{Read, Seek, SeekFrom, Write}; -use std::path::Path; -use std::time::Instant; - -struct Euclidean; - -impl Metric<[f32; 128]> for Euclidean { - type Unit = u32; - fn distance(&self, a: &[f32; 128], b: &[f32; 128]) -> u32 { - let sum: f32 = a - .iter() - .zip(b.iter()) - .map(|(&a, &b)| (a - b).powi(2)) - .sum(); - (sum.sqrt() * 1_000_000.0) as u32 - } -} - -const FEATURE_SIZE: usize = 128; -const FEATURE_BYTES: usize = FEATURE_SIZE * std::mem::size_of::(); - -/// A disk-based feature store that reads features from disk on every access. -/// This uses NO RAM for feature storage - only a small buffer for reading. -/// -/// Trade-off: Much slower due to disk I/O, but uses minimal memory. we are not -/// using mmapping here to isolate memory usage as low as possible (mmap is faster but -/// measurements are affected by OS page caching. realistically the results of this storage impl -/// show the true memory usage). -struct DiskFeatureStore { - file: File, - len: usize, -} - -impl DiskFeatureStore { - fn new>(path: P) -> std::io::Result { - let file = OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(true) - .open(path)?; - - Ok(Self { file, len: 0 }) - } -} - -impl FeatureStore<[f32; 128]> for DiskFeatureStore { - fn get(&self, index: usize) -> &[f32; 128] { - // Use thread-local storage to return a reference - // This is the key trick - we only keep ONE feature in memory at a time - thread_local! { - static BUFFER: std::cell::UnsafeCell<[f32; 128]> = std::cell::UnsafeCell::new([0.0; 128]); - } - - BUFFER.with(|buf| { - let buffer = unsafe { &mut *buf.get() }; - - let mut file = &self.file; - let offset = index * FEATURE_BYTES; - file.seek(SeekFrom::Start(offset as u64)).unwrap(); - - let mut bytes = [0u8; FEATURE_BYTES]; - file.read_exact(&mut bytes).unwrap(); - - for (i, chunk) in bytes.chunks_exact(4).enumerate() { - buffer[i] = f32::from_le_bytes(chunk.try_into().unwrap()); - } - - unsafe { &*(buffer as *const [f32; 128]) } - }) - } - - fn push(&mut self, feature: [f32; 128]) { - self.file.seek(SeekFrom::End(0)).unwrap(); - let bytes: Vec = feature.iter().flat_map(|f| f.to_le_bytes()).collect(); - self.file.write_all(&bytes).unwrap(); - self.len += 1; - } - - fn len(&self) -> usize { - self.len - } - - fn is_empty(&self) -> bool { - self.len == 0 - } -} - -/// A "null" feature store that doesn't store features at all. -/// Only useful for measuring the memory overhead of the graph structure alone. -struct NullFeatureStore { - len: usize, - dummy: [f32; 128], -} - -impl NullFeatureStore { - fn new() -> Self { - Self { - len: 0, - dummy: [0.0; 128], - } - } -} - -impl FeatureStore<[f32; 128]> for NullFeatureStore { - fn get(&self, _index: usize) -> &[f32; 128] { - // WARNING: This returns garbage! Only for memory measurement. - &self.dummy - } - - fn push(&mut self, _feature: [f32; 128]) { - self.len += 1; - } - - fn len(&self) -> usize { - self.len - } - - fn is_empty(&self) -> bool { - self.len == 0 - } -} - -fn generate_random_vectors(count: usize) -> Vec<[f32; 128]> { - let mut rng = rand::thread_rng(); - (0..count) - .map(|_| { - let mut v = [0.0f32; 128]; - for x in v.iter_mut() { - *x = rng.gen(); - } - v - }) - .collect() -} - -fn get_memory_usage() -> usize { - #[cfg(target_os = "linux")] - { - if let Ok(statm) = std::fs::read_to_string("/proc/self/statm") { - let parts: Vec<&str> = statm.split_whitespace().collect(); - if parts.len() >= 2 { - let rss_pages: usize = parts[1].parse().unwrap_or(0); - return rss_pages * 4096; - } - } - } - - #[cfg(target_os = "macos")] - { - use std::process::Command; - if let Ok(output) = Command::new("ps") - .args(&["-o", "rss=", "-p", &std::process::id().to_string()]) - .output() - { - if let Ok(rss_str) = String::from_utf8(output.stdout) { - if let Ok(rss_kb) = rss_str.trim().parse::() { - return rss_kb * 1024; - } - } - } - } - - 0 -} - -fn format_bytes(bytes: usize) -> String { - if bytes >= 1024 * 1024 * 1024 { - format!("{:.2} GB", bytes as f64 / (1024.0 * 1024.0 * 1024.0)) - } else if bytes >= 1024 * 1024 { - format!("{:.2} MB", bytes as f64 / (1024.0 * 1024.0)) - } else if bytes >= 1024 { - format!("{:.2} KB", bytes as f64 / 1024.0) - } else { - format!("{} bytes", bytes) - } -} - -fn main() -> std::io::Result<()> { - const NUM_VECTORS: usize = 100_000; - const NUM_QUERIES: usize = 10; - - let params = Params::new().ef_construction(20); // Lower for faster insertion - - // Generate test data - println!("Generating random vectors..."); - let vectors = generate_random_vectors(NUM_VECTORS); - let queries = generate_random_vectors(NUM_QUERIES); - - // Drop the generated vectors from consideration - let vectors_size = vectors.len() * FEATURE_BYTES; - - println!("=== Test 1: Graph structure only (NullFeatureStore) ==="); - - let mem_before_null = get_memory_usage(); - let graph_only_delta; - { - let storage = NullFeatureStore::new(); - let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); - - let mut hnsw: Hnsw = - Hnsw::new_with_storage_and_params(Euclidean, storage, params.clone(), prng); - let mut searcher = Searcher::default(); - - let start = Instant::now(); - for v in &vectors { - hnsw.insert(*v, &mut searcher); - } - - let mem_after = get_memory_usage(); - graph_only_delta = mem_after.saturating_sub(mem_before_null); - println!("Graph-only memory: +{}", format_bytes(graph_only_delta)); - } - println!(); - - println!("=== Test 2: In-memory Vec (default) ==="); - - let mem_before_vec = get_memory_usage(); - let vec_delta; - { - let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); - let mut hnsw: Hnsw = - Hnsw::new_params_and_prng(Euclidean, params.clone(), prng); - let mut searcher = Searcher::default(); - - let start = Instant::now(); - for v in &vectors { - hnsw.insert(*v, &mut searcher); - } - - let mem_after = get_memory_usage(); - vec_delta = mem_after.saturating_sub(mem_before_vec); - println!("Total memory: +{}", format_bytes(vec_delta)); - - let mut neighbors = [Neighbor { index: !0, distance: !0 }; 10]; - let start = Instant::now(); - for q in &queries { - hnsw.nearest(q, 64, &mut searcher, &mut neighbors); - } - println!("Search time: {:?} ({} queries)", start.elapsed(), NUM_QUERIES); - } - println!(); - - println!("=== Test 3: Disk-based FeatureStore ==="); - - let disk_path = "/tmp/hnsw_disk_features.bin"; - let mem_before_disk = get_memory_usage(); - let disk_delta; - { - let storage = DiskFeatureStore::new(disk_path)?; - let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); - - let mut hnsw: Hnsw = - Hnsw::new_with_storage_and_params(Euclidean, storage, params.clone(), prng); - let mut searcher = Searcher::default(); - - let start = Instant::now(); - for v in &vectors { - hnsw.insert(*v, &mut searcher); - } - - let mem_after = get_memory_usage(); - disk_delta = mem_after.saturating_sub(mem_before_disk); - - // Search (will be slow due to disk I/O) - let mut neighbors = [Neighbor { index: !0, distance: !0 }; 10]; - let start = Instant::now(); - for q in &queries { - hnsw.nearest(q, 64, &mut searcher, &mut neighbors); - } - println!("Search time: {:?} ({} queries) - slower due to disk I/O", start.elapsed(), NUM_QUERIES); - - std::fs::remove_file(disk_path).ok(); - } - println!(); - - println!("╔══════════════════════════════════════════════════════════════╗"); - println!("║ MEMORY SUMMARY ║"); - println!("╠══════════════════════════════════════════════════════════════╣"); - println!("║ Feature data size: {:>30} ║", format_bytes(NUM_VECTORS * FEATURE_BYTES)); - println!("╠══════════════════════════════════════════════════════════════╣"); - println!("║ Graph structure only: {:>30} ║", format_bytes(graph_only_delta)); - println!("║ Vec (graph + features): {:>30} ║", format_bytes(vec_delta)); - println!("║ Disk-based (graph only): {:>30} ║", format_bytes(disk_delta)); - println!("╠══════════════════════════════════════════════════════════════╣"); - - let feature_overhead = vec_delta.saturating_sub(graph_only_delta); - let savings = vec_delta.saturating_sub(disk_delta); - let savings_pct = if vec_delta > 0 { - (savings as f64 / vec_delta as f64) * 100.0 - } else { - 0.0 - }; - - println!("║ Feature storage overhead: {:>30} ║", format_bytes(feature_overhead)); - println!("║ Memory saved with disk: {:>30} ║", format_bytes(savings)); - println!("║ Savings percentage: {:>29.1}% ║", savings_pct); - println!("╚══════════════════════════════════════════════════════════════╝"); - println!(); - - Ok(()) -} diff --git a/examples/memory_usage.rs b/examples/memory_usage.rs new file mode 100644 index 0000000..b7eab65 --- /dev/null +++ b/examples/memory_usage.rs @@ -0,0 +1,231 @@ +use hnsw::{FeatureStore, Hnsw, Params, Searcher}; +use rand::Rng; +use rand_pcg::Pcg64; +use space::Metric; +use std::convert::TryInto; +use std::fs::{File, OpenOptions}; +use std::io::{Read, Seek, SeekFrom, Write}; +use std::path::Path; +use std::time::Instant; +use structopt::StructOpt; + +struct Euclidean; + +impl Metric<[f32; 128]> for Euclidean { + type Unit = u32; + fn distance(&self, a: &[f32; 128], b: &[f32; 128]) -> u32 { + let sum: f32 = a + .iter() + .zip(b.iter()) + .map(|(&a, &b)| (a - b).powi(2)) + .sum(); + (sum.sqrt() * 1_000_000.0) as u32 + } +} + +const FEATURE_SIZE: usize = 128; +const FEATURE_BYTES: usize = FEATURE_SIZE * std::mem::size_of::(); + +/// A disk-based feature store that reads features from disk on every access. +/// This uses NO RAM for feature storage - only a small buffer for reading. +struct DiskFeatureStore { + file: File, + len: usize, +} + +impl DiskFeatureStore { + fn new>(path: P) -> std::io::Result { + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(path)?; + + Ok(Self { file, len: 0 }) + } +} + +impl FeatureStore<[f32; 128]> for DiskFeatureStore { + fn get(&self, index: usize) -> &[f32; 128] { + thread_local! { + static BUFFER: std::cell::UnsafeCell<[f32; 128]> = std::cell::UnsafeCell::new([0.0; 128]); + } + + BUFFER.with(|buf| { + let buffer = unsafe { &mut *buf.get() }; + + let mut file = &self.file; + let offset = index * FEATURE_BYTES; + file.seek(SeekFrom::Start(offset as u64)).unwrap(); + + let mut bytes = [0u8; FEATURE_BYTES]; + file.read_exact(&mut bytes).unwrap(); + + for (i, chunk) in bytes.chunks_exact(4).enumerate() { + buffer[i] = f32::from_le_bytes(chunk.try_into().unwrap()); + } + + unsafe { &*(buffer as *const [f32; 128]) } + }) + } + + fn push(&mut self, feature: [f32; 128]) { + self.file.seek(SeekFrom::End(0)).unwrap(); + let bytes: Vec = feature.iter().flat_map(|f| f.to_le_bytes()).collect(); + self.file.write_all(&bytes).unwrap(); + self.len += 1; + } + + fn len(&self) -> usize { + self.len + } + + fn is_empty(&self) -> bool { + self.len == 0 + } +} + +#[derive(Debug, StructOpt)] +#[structopt( + name = "lazy_memory_comparison", + about = "Measures memory usage of HNSW with different storage backends" +)] +struct Opt { + /// Number of vectors to insert. + #[structopt(short = "n", long = "num-vectors", default_value = "100000")] + num_vectors: usize, + /// Use disk-based feature storage instead of in-memory Vec. + #[structopt(long = "disk")] + disk: bool, +} + +fn generate_random_vectors(count: usize) -> Vec<[f32; 128]> { + let mut rng = rand::thread_rng(); + (0..count) + .map(|_| { + let mut v = [0.0f32; 128]; + for x in v.iter_mut() { + *x = rng.gen(); + } + v + }) + .collect() +} + +fn get_memory_usage() -> usize { + #[cfg(target_os = "linux")] + { + if let Ok(statm) = std::fs::read_to_string("/proc/self/statm") { + let parts: Vec<&str> = statm.split_whitespace().collect(); + if parts.len() >= 2 { + let rss_pages: usize = parts[1].parse().unwrap_or(0); + return rss_pages * 4096; + } + } + } + + #[cfg(target_os = "macos")] + { + use std::process::Command; + if let Ok(output) = Command::new("ps") + .args(&["-o", "rss=", "-p", &std::process::id().to_string()]) + .output() + { + if let Ok(rss_str) = String::from_utf8(output.stdout) { + if let Ok(rss_kb) = rss_str.trim().parse::() { + return rss_kb * 1024; + } + } + } + } + + 0 +} + +fn format_bytes(bytes: usize) -> String { + if bytes >= 1024 * 1024 * 1024 { + format!("{:.2} GB", bytes as f64 / (1024.0 * 1024.0 * 1024.0)) + } else if bytes >= 1024 * 1024 { + format!("{:.2} MB", bytes as f64 / (1024.0 * 1024.0)) + } else if bytes >= 1024 { + format!("{:.2} KB", bytes as f64 / 1024.0) + } else { + format!("{} bytes", bytes) + } +} + +fn run_with_vec(num_vectors: usize) { + let params = Params::new().ef_construction(16); + + println!("=== Vec Storage ==="); + println!("Generating {} random vectors...", num_vectors); + let vectors = generate_random_vectors(num_vectors); + + let mem_before = get_memory_usage(); + + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + let mut hnsw: Hnsw = + Hnsw::new_params_and_prng(Euclidean, params, prng); + let mut searcher = Searcher::default(); + + println!("Building HNSW..."); + for v in &vectors { + hnsw.insert(*v, &mut searcher); + } + + let mem_after = get_memory_usage(); + let mem_delta = mem_after.saturating_sub(mem_before); + + println!(); + println!("Vectors: {}", num_vectors); + println!("Feature size: {} bytes", FEATURE_BYTES); + println!("Memory used: {}", format_bytes(mem_delta)); + println!("Expected features: {}", format_bytes(num_vectors * FEATURE_BYTES)); +} + +fn run_with_disk(num_vectors: usize) -> std::io::Result<()> { + let params = Params::new().ef_construction(16); + let disk_path = "/tmp/hnsw_disk_features.bin"; + + println!("=== Disk-based Storage ==="); + println!("Generating {} random vectors...", num_vectors); + let vectors = generate_random_vectors(num_vectors); + + let mem_before = get_memory_usage(); + + let storage = DiskFeatureStore::new(disk_path)?; + let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); + let mut hnsw: Hnsw = + Hnsw::new_with_storage_and_params(Euclidean, storage, params, prng); + let mut searcher = Searcher::default(); + + println!("Building HNSW..."); + for v in &vectors { + hnsw.insert(*v, &mut searcher); + } + + let mem_after = get_memory_usage(); + let mem_delta = mem_after.saturating_sub(mem_before); + + println!(); + println!("Vectors: {}", num_vectors); + println!("Feature size: {} bytes", FEATURE_BYTES); + println!("Memory used: {}", format_bytes(mem_delta)); + println!("Expected features: {}", format_bytes(num_vectors * FEATURE_BYTES)); + + std::fs::remove_file(disk_path).ok(); + Ok(()) +} + +fn main() -> std::io::Result<()> { + let opt = Opt::from_args(); + + if opt.disk { + run_with_disk(opt.num_vectors)?; + } else { + run_with_vec(opt.num_vectors); + } + + Ok(()) +} From 9bdf45f12ac08b540d7002cd86a3adbe915bc577 Mon Sep 17 00:00:00 2001 From: Giulio Rebuffo Date: Thu, 11 Dec 2025 19:46:32 +0100 Subject: [PATCH 3/8] save --- examples/recall_discrete.rs | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/examples/recall_discrete.rs b/examples/recall_discrete.rs index 9afb76e..1695c7d 100644 --- a/examples/recall_discrete.rs +++ b/examples/recall_discrete.rs @@ -14,14 +14,22 @@ use std::io::Read; use std::path::PathBuf; use structopt::StructOpt; -/// A mmap-based feature store for BitArray types. + struct MmapBitArrayStore { mmap: MmapMut, len: usize, capacity: usize, + stride: usize, _marker: std::marker::PhantomData, } +const ALIGNMENT: usize = 64; + +// this is very suboptimal but works for testing purposes +fn align_up(size: usize, align: usize) -> usize { + (size + align - 1) & !(align - 1) +} + impl MmapBitArrayStore { fn new(path: &str, capacity: usize, item_size: usize) -> std::io::Result { let file = OpenOptions::new() @@ -31,7 +39,10 @@ impl MmapBitArrayStore { .truncate(true) .open(path)?; - let total_size = capacity * item_size; + // Align each item to 64 bytes for SIMD + let stride = align_up(item_size, ALIGNMENT); + // Add padding at start to ensure first item is aligned + let total_size = ALIGNMENT + capacity * stride; file.set_len(total_size as u64)?; let mmap = unsafe { MmapMut::map_mut(&file)? }; @@ -40,21 +51,30 @@ impl MmapBitArrayStore { mmap, len: 0, capacity, + stride, _marker: std::marker::PhantomData, }) } + + fn aligned_offset(&self, index: usize) -> usize { + // Find offset that gives 64-byte alignment + let base = self.mmap.as_ptr() as usize; + let aligned_base = align_up(base, ALIGNMENT); + let padding = aligned_base - base; + padding + index * self.stride + } } impl FeatureStore> for MmapBitArrayStore> { fn get(&self, index: usize) -> &BitArray { - let offset = index * N; - // Safety: BitArray is repr(transparent) over [u8; N] + let offset = self.aligned_offset(index); unsafe { &*(self.mmap[offset..].as_ptr() as *const BitArray) } } fn push(&mut self, feature: BitArray) { + // we dont need to be sophisticated for this dummy example assert!(self.len < self.capacity, "MmapBitArrayStore capacity exceeded"); - let offset = self.len * N; + let offset = self.aligned_offset(self.len); self.mmap[offset..offset + N].copy_from_slice(feature.bytes()); self.len += 1; } From 8fe886b396e752c3a7f3d09ba3f8993cb50d0873 Mon Sep 17 00:00:00 2001 From: Giulio Rebuffo Date: Thu, 11 Dec 2025 19:59:19 +0100 Subject: [PATCH 4/8] save --- examples/memory_usage.rs | 14 --- examples/recall.rs | 221 ++++++++++++++++++++++++++------------- 2 files changed, 149 insertions(+), 86 deletions(-) diff --git a/examples/memory_usage.rs b/examples/memory_usage.rs index b7eab65..90412b5 100644 --- a/examples/memory_usage.rs +++ b/examples/memory_usage.rs @@ -9,20 +9,6 @@ use std::path::Path; use std::time::Instant; use structopt::StructOpt; -struct Euclidean; - -impl Metric<[f32; 128]> for Euclidean { - type Unit = u32; - fn distance(&self, a: &[f32; 128], b: &[f32; 128]) -> u32 { - let sum: f32 = a - .iter() - .zip(b.iter()) - .map(|(&a, &b)| (a - b).powi(2)) - .sum(); - (sum.sqrt() * 1_000_000.0) as u32 - } -} - const FEATURE_SIZE: usize = 128; const FEATURE_BYTES: usize = FEATURE_SIZE * std::mem::size_of::(); diff --git a/examples/recall.rs b/examples/recall.rs index 3e8b46b..d2383fc 100644 --- a/examples/recall.rs +++ b/examples/recall.rs @@ -97,15 +97,33 @@ impl FeatureStore<[f32; 64]> for MmapFeatureStore { #[structopt(name = "recall", about = "Generates recall graphs for HNSW")] struct Opt { /// The value of M to use. + /// + /// This can only be between 4 and 52 inclusive and a multiple of 4. + /// M0 is set to 2 * M. #[structopt(short = "m", long = "max_edges", default_value = "24")] m: usize, /// The dataset size to test on. #[structopt(short = "s", long = "size", default_value = "10000")] size: usize, /// Total number of query bitstrings. + /// + /// The higher this is, the better the quality of the output data and statistics, but + /// the longer the benchmark will take to set up. #[structopt(short = "q", long = "queries", default_value = "10000")] num_queries: usize, /// The number of dimensions in the feature vector. + /// + /// This is the length of the feature vector. The descriptor_stride (-d) + /// parameter must exceed this value. + /// + /// Possible values: + /// - 8 + /// - 16 + /// - 32 + /// - 64 + /// - 128 + /// - 256 + /// - 512 #[structopt(short = "l", long = "dimensions", default_value = "64")] dimensions: usize, /// The beginning ef value. @@ -121,14 +139,18 @@ struct Opt { #[structopt(short = "f", long = "file")] file: Option, /// The descriptor stride length in floats. + /// + /// KAZE: 64 + /// SIFT: 128 #[structopt(short = "d", long = "descriptor_stride", default_value = "64")] descriptor_stride: usize, /// efConstruction controlls the quality of the graph at build-time. #[structopt(short = "c", long = "ef_construction", default_value = "400")] ef_construction: usize, /// Use mmap-based feature storage instead of in-memory Vec. - #[structopt(long = "mmap")] - mmap: bool, + /// Note: Only supports 64 dimensions and does not support file input. + #[structopt(long = "disk")] + disk: bool, } fn process(opt: &Opt) -> (Vec, Vec) { @@ -141,37 +163,54 @@ fn process(opt: &Opt) -> (Vec, Vec) { let (search_space, query_strings): (Vec, Vec) = if let Some(filepath) = &opt.file { eprintln!( "Reading {} search space descriptors of size {} f32s from file \"{}\"...", - opt.size, opt.descriptor_stride, filepath.display() + opt.size, + opt.descriptor_stride, + filepath.display() ); let mut file = std::fs::File::open(filepath).expect("unable to open file"); + // We are loading floats, so multiply by 4. let mut search_space = vec![0u8; opt.size * opt.descriptor_stride * 4]; file.read_exact(&mut search_space).expect( "unable to read enough search descriptors from the file (try decreasing -s/-q)", ); - let search_space = search_space.chunks_exact(4).map(LittleEndian::read_f32).collect(); + let search_space = search_space + .chunks_exact(4) + .map(LittleEndian::read_f32) + .collect(); eprintln!("Done."); eprintln!( "Reading {} query descriptors of size {} f32s from file \"{}\"...", - opt.num_queries, opt.descriptor_stride, filepath.display() + opt.num_queries, + opt.descriptor_stride, + filepath.display() ); + // We are loading floats, so multiply by 4. let mut query_strings = vec![0u8; opt.num_queries * opt.descriptor_stride * 4]; file.read_exact(&mut query_strings) .expect("unable to read enough query descriptors from the file (try decreasing -q/-s)"); - let query_strings = query_strings.chunks_exact(4).map(LittleEndian::read_f32).collect(); + let query_strings = query_strings + .chunks_exact(4) + .map(LittleEndian::read_f32) + .collect(); eprintln!("Done."); (search_space, query_strings) } else { - eprintln!("Generating {} random vectors...", opt.size); + eprintln!("Generating {} random bitstrings...", opt.size); let search_space: Vec = rng .sample_iter(&Standard) .take(opt.size * opt.descriptor_stride) .collect(); eprintln!("Done."); + // Create another RNG to prevent potential correlation. let rng = Pcg64::from_seed([6; 32]); - eprintln!("Generating {} independent random query vectors...", opt.num_queries); + + eprintln!( + "Generating {} independent random query strings...", + opt.num_queries + ); let query_strings: Vec = rng .sample_iter(&Standard) .take(opt.num_queries * opt.descriptor_stride) @@ -189,7 +228,10 @@ fn process(opt: &Opt) -> (Vec, Vec) { .map(|c| &c[..opt.dimensions]) .collect(); - eprintln!("Computing the correct nearest neighbor distance for all {} queries...", opt.num_queries); + eprintln!( + "Computing the correct nearest neighbor distance for all {} queries...", + opt.num_queries + ); let correct_worst_distances: Vec<_> = query_strings .iter() .cloned() @@ -202,6 +244,7 @@ fn process(opt: &Opt) -> (Vec, Vec) { v.resize_with(opt.k, || unreachable!()); } } + // Get the worst distance v.into_iter().take(opt.k).last().unwrap() }) .collect(); @@ -224,13 +267,21 @@ fn process(opt: &Opt) -> (Vec, Vec) { let (recalls, times): (Vec, Vec) = efs .map(|ef| { let correct = RefCell::new(0usize); - let dest = vec![Neighbor { index: !0, distance: !0 }; opt.k]; + let dest = vec![ + Neighbor { + index: !0, + distance: !0, + }; + opt.k + ]; let stats = easybench::bench_env(dest, |mut dest| { let mut refmut = state.borrow_mut(); let (searcher, query) = &mut *refmut; let (ix, query_feature) = query.next().unwrap(); let correct_worst_distance = correct_worst_distances[ix]; + // Go through all the features. for &mut neighbor in hnsw.nearest(&query_feature, ef, searcher, &mut dest) { + // Any feature that is less than or equal to the worst real nearest neighbor distance is correct. if Euclidean.distance(&search_space[neighbor.index], &query_feature) <= correct_worst_distance { @@ -240,11 +291,15 @@ fn process(opt: &Opt) -> (Vec, Vec) { }); (stats, correct.into_inner()) }) - .fold((vec![], vec![]), |(mut recalls, mut times), (stats, correct)| { - times.push((stats.ns_per_iter * 0.1f64.powi(9)).recip()); - recalls.push(correct as f64 / (stats.iterations * opt.k) as f64); - (recalls, times) - }); + .fold( + (vec![], vec![]), + |(mut recalls, mut times), (stats, correct)| { + times.push((stats.ns_per_iter * 0.1f64.powi(9)).recip()); + // The maximum number of correct nearest neighbors is + recalls.push(correct as f64 / (stats.iterations * opt.k) as f64); + (recalls, times) + }, + ); eprintln!("Done."); (recalls, times) @@ -254,15 +309,21 @@ fn process_mmap, const M: usize, const M0: usize>( opt: &Opt, storage: S, ) -> (Vec, Vec) { - assert!(opt.k <= opt.size, "You must choose a dataset size larger or equal to the test search size"); - assert!(opt.dimensions == 64, "Mmap mode only supports 64 dimensions (use -l 64)"); + assert!( + opt.k <= opt.size, + "You must choose a dataset size larger or equal to the test search size" + ); + assert!(opt.dimensions == 64, "Mmap mode only supports 64 dimensions"); assert!(opt.file.is_none(), "Mmap mode does not support file input"); let rng = Pcg64::from_seed([5; 32]); eprintln!("Generating {} random vectors...", opt.size); - let raw: Vec = rng.clone().sample_iter(&Standard).take(opt.size * 64).collect(); - let search_space: Vec<[f32; 64]> = raw + let search_space: Vec<[f32; 64]> = rng + .clone() + .sample_iter(&Standard) + .take(opt.size * 64) + .collect::>() .chunks_exact(64) .map(|chunk| { let mut arr = [0.0f32; 64]; @@ -272,10 +333,17 @@ fn process_mmap, const M: usize, const M0: usize>( .collect(); eprintln!("Done."); + // Create another RNG to prevent potential correlation. let rng = Pcg64::from_seed([6; 32]); - eprintln!("Generating {} independent random query vectors...", opt.num_queries); - let raw: Vec = rng.sample_iter(&Standard).take(opt.num_queries * 64).collect(); - let query_strings: Vec<[f32; 64]> = raw + + eprintln!( + "Generating {} independent random query strings...", + opt.num_queries + ); + let query_strings: Vec<[f32; 64]> = rng + .sample_iter(&Standard) + .take(opt.num_queries * 64) + .collect::>() .chunks_exact(64) .map(|chunk| { let mut arr = [0.0f32; 64]; @@ -285,7 +353,10 @@ fn process_mmap, const M: usize, const M0: usize>( .collect(); eprintln!("Done."); - eprintln!("Computing the correct nearest neighbor distance for all {} queries...", opt.num_queries); + eprintln!( + "Computing the correct nearest neighbor distance for all {} queries...", + opt.num_queries + ); let correct_worst_distances: Vec<_> = query_strings .iter() .map(|feature| { @@ -297,15 +368,20 @@ fn process_mmap, const M: usize, const M0: usize>( v.resize_with(opt.k, || unreachable!()); } } + // Get the worst distance v.into_iter().take(opt.k).last().unwrap() }) .collect(); eprintln!("Done."); - eprintln!("Generating HNSW with MmapFeatureStore..."); + eprintln!("Generating HNSW..."); let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); - let mut hnsw: Hnsw = - Hnsw::new_with_storage_and_params(Euclidean, storage, Params::new().ef_construction(opt.ef_construction), prng); + let mut hnsw: Hnsw = Hnsw::new_with_storage_and_params( + Euclidean, + storage, + Params::new().ef_construction(opt.ef_construction), + prng, + ); let mut searcher: Searcher<_> = Searcher::default(); for feature in &search_space { hnsw.insert(*feature, &mut searcher); @@ -318,13 +394,21 @@ fn process_mmap, const M: usize, const M0: usize>( let (recalls, times): (Vec, Vec) = efs .map(|ef| { let correct = RefCell::new(0usize); - let dest = vec![Neighbor { index: !0, distance: !0 }; opt.k]; + let dest = vec![ + Neighbor { + index: !0, + distance: !0, + }; + opt.k + ]; let stats = easybench::bench_env(dest, |mut dest| { let mut refmut = state.borrow_mut(); let (searcher, query) = &mut *refmut; let (ix, query_feature) = query.next().unwrap(); let correct_worst_distance = correct_worst_distances[ix]; + // Go through all the features. for &mut neighbor in hnsw.nearest(&query_feature, ef, searcher, &mut dest) { + // Any feature that is less than or equal to the worst real nearest neighbor distance is correct. if Euclidean.distance(&search_space[neighbor.index], &query_feature) <= correct_worst_distance { @@ -334,46 +418,38 @@ fn process_mmap, const M: usize, const M0: usize>( }); (stats, correct.into_inner()) }) - .fold((vec![], vec![]), |(mut recalls, mut times), (stats, correct)| { - times.push((stats.ns_per_iter * 0.1f64.powi(9)).recip()); - recalls.push(correct as f64 / (stats.iterations * opt.k) as f64); - (recalls, times) - }); + .fold( + (vec![], vec![]), + |(mut recalls, mut times), (stats, correct)| { + times.push((stats.ns_per_iter * 0.1f64.powi(9)).recip()); + // The maximum number of correct nearest neighbors is + recalls.push(correct as f64 / (stats.iterations * opt.k) as f64); + (recalls, times) + }, + ); eprintln!("Done."); (recalls, times) } -macro_rules! process_m { - ($opt:expr, $m:expr, $m0:expr) => { - process::<$m, $m0>(&$opt) - }; -} - -macro_rules! process_mmap_m { - ($opt:expr, $m:expr, $m0:expr) => { - process_mmap::<_, $m, $m0>(&$opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", $opt.size).unwrap()) - }; -} - fn main() { let opt = Opt::from_args(); - let (recalls, times, storage_type) = if opt.mmap { + let (recalls, times, storage_type) = if opt.disk { let (r, t) = match opt.m { - 4 => process_mmap_m!(opt, 4, 8), - 8 => process_mmap_m!(opt, 8, 16), - 12 => process_mmap_m!(opt, 12, 24), - 16 => process_mmap_m!(opt, 16, 32), - 20 => process_mmap_m!(opt, 20, 40), - 24 => process_mmap_m!(opt, 24, 48), - 28 => process_mmap_m!(opt, 28, 56), - 32 => process_mmap_m!(opt, 32, 64), - 36 => process_mmap_m!(opt, 36, 72), - 40 => process_mmap_m!(opt, 40, 80), - 44 => process_mmap_m!(opt, 44, 88), - 48 => process_mmap_m!(opt, 48, 96), - 52 => process_mmap_m!(opt, 52, 104), + 4 => process_mmap::<_, 4, 8>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 8 => process_mmap::<_, 8, 16>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 12 => process_mmap::<_, 12, 24>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 16 => process_mmap::<_, 16, 32>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 20 => process_mmap::<_, 20, 40>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 24 => process_mmap::<_, 24, 48>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 28 => process_mmap::<_, 28, 56>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 32 => process_mmap::<_, 32, 64>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 36 => process_mmap::<_, 36, 72>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 40 => process_mmap::<_, 40, 80>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 44 => process_mmap::<_, 44, 88>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 48 => process_mmap::<_, 48, 96>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 52 => process_mmap::<_, 52, 104>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), _ => { eprintln!("Only M between 4 and 52 inclusive and multiples of 4 are allowed"); return; @@ -382,19 +458,19 @@ fn main() { (r, t, "MmapFeatureStore") } else { let (r, t) = match opt.m { - 4 => process_m!(opt, 4, 8), - 8 => process_m!(opt, 8, 16), - 12 => process_m!(opt, 12, 24), - 16 => process_m!(opt, 16, 32), - 20 => process_m!(opt, 20, 40), - 24 => process_m!(opt, 24, 48), - 28 => process_m!(opt, 28, 56), - 32 => process_m!(opt, 32, 64), - 36 => process_m!(opt, 36, 72), - 40 => process_m!(opt, 40, 80), - 44 => process_m!(opt, 44, 88), - 48 => process_m!(opt, 48, 96), - 52 => process_m!(opt, 52, 104), + 4 => process::<4, 8>(&opt), + 8 => process::<8, 16>(&opt), + 12 => process::<12, 24>(&opt), + 16 => process::<16, 32>(&opt), + 20 => process::<20, 40>(&opt), + 24 => process::<24, 48>(&opt), + 28 => process::<28, 56>(&opt), + 32 => process::<32, 64>(&opt), + 36 => process::<36, 72>(&opt), + 40 => process::<40, 80>(&opt), + 44 => process::<44, 88>(&opt), + 48 => process::<48, 96>(&opt), + 52 => process::<52, 104>(&opt), _ => { eprintln!("Only M between 4 and 52 inclusive and multiples of 4 are allowed"); return; @@ -404,6 +480,7 @@ fn main() { }; let mut fg = Figure::new(); + fg.axes2d() .set_title( &format!( @@ -422,5 +499,5 @@ fn main() { .set_y_grid(true) .set_y_minor_grid(true); - fg.show().ok(); + fg.show().expect("unable to show gnuplot"); } From b166cf10c2cd58d2bcec19a855aa4ae3ebe22faf Mon Sep 17 00:00:00 2001 From: Giulio Rebuffo Date: Thu, 11 Dec 2025 20:03:23 +0100 Subject: [PATCH 5/8] save --- examples/recall.rs | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/examples/recall.rs b/examples/recall.rs index d2383fc..a82af69 100644 --- a/examples/recall.rs +++ b/examples/recall.rs @@ -305,7 +305,7 @@ fn process(opt: &Opt) -> (Vec, Vec) { (recalls, times) } -fn process_mmap, const M: usize, const M0: usize>( +fn process_disk, const M: usize, const M0: usize>( opt: &Opt, storage: S, ) -> (Vec, Vec) { @@ -313,12 +313,12 @@ fn process_mmap, const M: usize, const M0: usize>( opt.k <= opt.size, "You must choose a dataset size larger or equal to the test search size" ); - assert!(opt.dimensions == 64, "Mmap mode only supports 64 dimensions"); - assert!(opt.file.is_none(), "Mmap mode does not support file input"); + assert!(opt.dimensions == 64, "Disk mode only supports 64 dimensions"); + assert!(opt.file.is_none(), "Disk mode does not support file input"); let rng = Pcg64::from_seed([5; 32]); - eprintln!("Generating {} random vectors...", opt.size); + eprintln!("Generating {} random bitstrings...", opt.size); let search_space: Vec<[f32; 64]> = rng .clone() .sample_iter(&Standard) @@ -437,19 +437,19 @@ fn main() { let (recalls, times, storage_type) = if opt.disk { let (r, t) = match opt.m { - 4 => process_mmap::<_, 4, 8>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 8 => process_mmap::<_, 8, 16>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 12 => process_mmap::<_, 12, 24>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 16 => process_mmap::<_, 16, 32>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 20 => process_mmap::<_, 20, 40>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 24 => process_mmap::<_, 24, 48>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 28 => process_mmap::<_, 28, 56>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 32 => process_mmap::<_, 32, 64>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 36 => process_mmap::<_, 36, 72>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 40 => process_mmap::<_, 40, 80>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 44 => process_mmap::<_, 44, 88>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 48 => process_mmap::<_, 48, 96>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 52 => process_mmap::<_, 52, 104>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 4 => process_disk::<_, 4, 8>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 8 => process_disk::<_, 8, 16>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 12 => process_disk::<_, 12, 24>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 16 => process_disk::<_, 16, 32>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 20 => process_disk::<_, 20, 40>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 24 => process_disk::<_, 24, 48>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 28 => process_disk::<_, 28, 56>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 32 => process_disk::<_, 32, 64>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 36 => process_disk::<_, 36, 72>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 40 => process_disk::<_, 40, 80>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 44 => process_disk::<_, 44, 88>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 48 => process_disk::<_, 48, 96>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 52 => process_disk::<_, 52, 104>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), _ => { eprintln!("Only M between 4 and 52 inclusive and multiples of 4 are allowed"); return; From 4b372e8a16f19391ff6bc72f815d817f4825907c Mon Sep 17 00:00:00 2001 From: Giulio Rebuffo Date: Thu, 11 Dec 2025 20:14:03 +0100 Subject: [PATCH 6/8] save --- examples/recall.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/recall.rs b/examples/recall.rs index a82af69..0d35912 100644 --- a/examples/recall.rs +++ b/examples/recall.rs @@ -148,7 +148,6 @@ struct Opt { #[structopt(short = "c", long = "ef_construction", default_value = "400")] ef_construction: usize, /// Use mmap-based feature storage instead of in-memory Vec. - /// Note: Only supports 64 dimensions and does not support file input. #[structopt(long = "disk")] disk: bool, } From e61693e7c791adb3de9966bf10284a50c9678b98 Mon Sep 17 00:00:00 2001 From: Giulio Rebuffo Date: Thu, 11 Dec 2025 20:20:03 +0100 Subject: [PATCH 7/8] save --- examples/recall.rs | 98 +++++++++++++++++++++++++++++----------------- 1 file changed, 61 insertions(+), 37 deletions(-) diff --git a/examples/recall.rs b/examples/recall.rs index 0d35912..9b09e81 100644 --- a/examples/recall.rs +++ b/examples/recall.rs @@ -27,9 +27,9 @@ impl Metric<&[f32]> for Euclidean { } } -impl Metric<[f32; 64]> for Euclidean { +impl Metric<[f32; N]> for Euclidean { type Unit = u32; - fn distance(&self, a: &[f32; 64], b: &[f32; 64]) -> u32 { + fn distance(&self, a: &[f32; N], b: &[f32; N]) -> u32 { a.iter() .zip(b.iter()) .map(|(&a, &b)| (a - b).powi(2)) @@ -39,14 +39,14 @@ impl Metric<[f32; 64]> for Euclidean { } } -/// A mmap-based feature store for [f32; 64] arrays. -struct MmapFeatureStore { +/// A mmap-based feature store for [f32; N] arrays. +struct MmapFeatureStore { mmap: MmapMut, len: usize, capacity: usize, } -impl MmapFeatureStore { +impl MmapFeatureStore { fn new(path: &str, capacity: usize) -> std::io::Result { let file = OpenOptions::new() .read(true) @@ -55,7 +55,7 @@ impl MmapFeatureStore { .truncate(true) .open(path)?; - let total_size = capacity * std::mem::size_of::<[f32; 64]>(); + let total_size = capacity * std::mem::size_of::<[f32; N]>(); file.set_len(total_size as u64)?; let mmap = unsafe { MmapMut::map_mut(&file)? }; @@ -68,17 +68,17 @@ impl MmapFeatureStore { } } -impl FeatureStore<[f32; 64]> for MmapFeatureStore { - fn get(&self, index: usize) -> &[f32; 64] { - let offset = index * std::mem::size_of::<[f32; 64]>(); - unsafe { &*(self.mmap[offset..].as_ptr() as *const [f32; 64]) } +impl FeatureStore<[f32; N]> for MmapFeatureStore { + fn get(&self, index: usize) -> &[f32; N] { + let offset = index * std::mem::size_of::<[f32; N]>(); + unsafe { &*(self.mmap[offset..].as_ptr() as *const [f32; N]) } } - fn push(&mut self, feature: [f32; 64]) { + fn push(&mut self, feature: [f32; N]) { assert!(self.len < self.capacity, "MmapFeatureStore capacity exceeded"); - let offset = self.len * std::mem::size_of::<[f32; 64]>(); + let offset = self.len * std::mem::size_of::<[f32; N]>(); let bytes: &[u8] = unsafe { - std::slice::from_raw_parts(feature.as_ptr() as *const u8, std::mem::size_of::<[f32; 64]>()) + std::slice::from_raw_parts(feature.as_ptr() as *const u8, std::mem::size_of::<[f32; N]>()) }; self.mmap[offset..offset + bytes.len()].copy_from_slice(bytes); self.len += 1; @@ -304,7 +304,8 @@ fn process(opt: &Opt) -> (Vec, Vec) { (recalls, times) } -fn process_disk, const M: usize, const M0: usize>( + +fn process_disk, const M: usize, const M0: usize, const N: usize>( opt: &Opt, storage: S, ) -> (Vec, Vec) { @@ -312,20 +313,19 @@ fn process_disk, const M: usize, const M0: usize>( opt.k <= opt.size, "You must choose a dataset size larger or equal to the test search size" ); - assert!(opt.dimensions == 64, "Disk mode only supports 64 dimensions"); assert!(opt.file.is_none(), "Disk mode does not support file input"); let rng = Pcg64::from_seed([5; 32]); eprintln!("Generating {} random bitstrings...", opt.size); - let search_space: Vec<[f32; 64]> = rng + let search_space: Vec<[f32; N]> = rng .clone() .sample_iter(&Standard) - .take(opt.size * 64) + .take(opt.size * N) .collect::>() - .chunks_exact(64) + .chunks_exact(N) .map(|chunk| { - let mut arr = [0.0f32; 64]; + let mut arr = [0.0f32; N]; arr.copy_from_slice(chunk); arr }) @@ -339,13 +339,13 @@ fn process_disk, const M: usize, const M0: usize>( "Generating {} independent random query strings...", opt.num_queries ); - let query_strings: Vec<[f32; 64]> = rng + let query_strings: Vec<[f32; N]> = rng .sample_iter(&Standard) - .take(opt.num_queries * 64) + .take(opt.num_queries * N) .collect::>() - .chunks_exact(64) + .chunks_exact(N) .map(|chunk| { - let mut arr = [0.0f32; 64]; + let mut arr = [0.0f32; N]; arr.copy_from_slice(chunk); arr }) @@ -375,7 +375,7 @@ fn process_disk, const M: usize, const M0: usize>( eprintln!("Generating HNSW..."); let prng = Pcg64::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7ac28fa16a64abf96); - let mut hnsw: Hnsw = Hnsw::new_with_storage_and_params( + let mut hnsw: Hnsw = Hnsw::new_with_storage_and_params( Euclidean, storage, Params::new().ef_construction(opt.ef_construction), @@ -431,24 +431,48 @@ fn process_disk, const M: usize, const M0: usize>( (recalls, times) } +macro_rules! process_disk_m { + ( $opt:expr, $m:expr, $m0:expr ) => { + match $opt.dimensions { + 64 => process_disk::<_, $m, $m0, 64>( + &$opt, + MmapFeatureStore::<64>::new("/tmp/recall_mmap.bin", $opt.size).unwrap(), + ), + 128 => process_disk::<_, $m, $m0, 128>( + &$opt, + MmapFeatureStore::<128>::new("/tmp/recall_mmap.bin", $opt.size).unwrap(), + ), + 256 => process_disk::<_, $m, $m0, 256>( + &$opt, + MmapFeatureStore::<256>::new("/tmp/recall_mmap.bin", $opt.size).unwrap(), + ), + 512 => process_disk::<_, $m, $m0, 512>( + &$opt, + MmapFeatureStore::<512>::new("/tmp/recall_mmap.bin", $opt.size).unwrap(), + ), + _ => panic!("error: incorrect dimensions for disk mode, supported: 64, 128, 256, 512"), + } + }; +} + fn main() { let opt = Opt::from_args(); let (recalls, times, storage_type) = if opt.disk { let (r, t) = match opt.m { - 4 => process_disk::<_, 4, 8>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 8 => process_disk::<_, 8, 16>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 12 => process_disk::<_, 12, 24>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 16 => process_disk::<_, 16, 32>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 20 => process_disk::<_, 20, 40>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 24 => process_disk::<_, 24, 48>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 28 => process_disk::<_, 28, 56>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 32 => process_disk::<_, 32, 64>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 36 => process_disk::<_, 36, 72>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 40 => process_disk::<_, 40, 80>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 44 => process_disk::<_, 44, 88>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 48 => process_disk::<_, 48, 96>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), - 52 => process_disk::<_, 52, 104>(&opt, MmapFeatureStore::new("/tmp/recall_mmap.bin", opt.size).unwrap()), + 4 => process_disk_m!(opt, 4, 8), + 8 => process_disk_m!(opt, 8, 16), + 12 => process_disk_m!(opt, 12, 24), + 16 => process_disk_m!(opt, 16, 32), + 20 => process_disk_m!(opt, 20, 40), + 24 => process_disk_m!(opt, 24, 48), + 28 => process_disk_m!(opt, 28, 56), + 32 => process_disk_m!(opt, 32, 64), + 36 => process_disk_m!(opt, 36, 72), + 40 => process_disk_m!(opt, 40, 80), + 44 => process_disk_m!(opt, 44, 88), + 48 => process_disk_m!(opt, 48, 96), + 52 => process_disk_m!(opt, 52, 104), _ => { eprintln!("Only M between 4 and 52 inclusive and multiples of 4 are allowed"); return; From fdf6caf80daac7a2bc9603615f8fc6c1a0e5e6be Mon Sep 17 00:00:00 2001 From: Giulio Rebuffo Date: Thu, 11 Dec 2025 20:30:58 +0100 Subject: [PATCH 8/8] ops --- examples/memory_usage.rs | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/examples/memory_usage.rs b/examples/memory_usage.rs index 90412b5..d4060b3 100644 --- a/examples/memory_usage.rs +++ b/examples/memory_usage.rs @@ -6,9 +6,22 @@ use std::convert::TryInto; use std::fs::{File, OpenOptions}; use std::io::{Read, Seek, SeekFrom, Write}; use std::path::Path; -use std::time::Instant; use structopt::StructOpt; +struct Euclidean; + +impl Metric<[f32; 128]> for Euclidean { + type Unit = u32; + fn distance(&self, a: &[f32; 128], b: &[f32; 128]) -> u32 { + a.iter() + .zip(b.iter()) + .map(|(&a, &b)| (a - b).powi(2)) + .sum::() + .sqrt() + .to_bits() + } +} + const FEATURE_SIZE: usize = 128; const FEATURE_BYTES: usize = FEATURE_SIZE * std::mem::size_of::();