use hashbrown::HashMap;
use std::cmp::Ordering;
use crate::hnsw;
use anndists::dist::distances::Distance;
use hnsw::*;
use log::error;
impl PartialEq for Neighbour {
fn eq(&self, other: &Neighbour) -> bool {
self.distance == other.distance
} }
impl Eq for Neighbour {}
#[allow(clippy::non_canonical_partial_ord_impl)]
impl PartialOrd for Neighbour {
fn partial_cmp(&self, other: &Neighbour) -> Option<Ordering> {
self.distance.partial_cmp(&other.distance)
} }
impl Ord for Neighbour {
fn cmp(&self, other: &Neighbour) -> Ordering {
if !self.distance.is_nan() && !other.distance.is_nan() {
self.distance.partial_cmp(&other.distance).unwrap()
} else {
panic!("got a NaN in a distance");
}
} }
#[derive(Clone)]
pub struct FlatPoint {
origin_id: DataId,
p_id: PointId,
neighbours: Vec<Neighbour>,
}
impl FlatPoint {
pub fn get_neighbours(&self) -> &Vec<Neighbour> {
&self.neighbours
}
pub fn get_id(&self) -> DataId {
self.origin_id
}
pub fn get_p_id(&self) -> PointId {
self.p_id
}
}
fn flatten_point<T: Clone + Send + Sync>(point: &Point<T>) -> FlatPoint {
let neighbours = point.get_neighborhood_id();
let mut flat_neighbours = Vec::<Neighbour>::new();
for layer in neighbours {
for neighbour in layer {
flat_neighbours.push(neighbour);
}
}
flat_neighbours.sort_unstable();
FlatPoint {
origin_id: point.get_origin_id(),
p_id: point.get_point_id(),
neighbours: flat_neighbours,
}
}
pub struct FlatNeighborhood {
hash_t: HashMap<DataId, FlatPoint>,
}
impl FlatNeighborhood {
pub fn get_neighbours(&self, p_id: DataId) -> Option<Vec<Neighbour>> {
self.hash_t
.get(&p_id)
.map(|point| point.get_neighbours().clone())
}
}
impl<T: Clone + Send + Sync, D: Distance<T> + Send + Sync> From<&Hnsw<'_, T, D>>
for FlatNeighborhood
{
fn from(hnsw: &Hnsw<T, D>) -> Self {
let mut hash_t = HashMap::new();
let pt_iter = hnsw.get_point_indexation().into_iter();
for point in pt_iter {
let res_insert = hash_t.insert(point.get_origin_id(), flatten_point(&point));
if let Some(old_point) = res_insert {
error!("2 points with same origin id {:?}", old_point.origin_id);
}
}
FlatNeighborhood { hash_t }
}
}
#[cfg(test)]
mod tests {
use super::*;
use anndists::dist::distances::*;
use log::debug;
use crate::api::AnnT;
use crate::hnswio::*;
use rand::distr::{Distribution, Uniform};
fn log_init_test() {
let _ = env_logger::builder().is_test(true).try_init();
}
#[test]
fn test_dump_reload_graph_flatten() {
println!("\n\n test_dump_reload_graph_flatten");
log_init_test();
let mut rng = rand::rng();
let unif = Uniform::<f32>::new(0., 1.).unwrap();
let nbcolumn = 1000;
let nbrow = 10;
let mut xsi;
let mut data = Vec::with_capacity(nbcolumn);
for j in 0..nbcolumn {
data.push(Vec::with_capacity(nbrow));
for _ in 0..nbrow {
xsi = unif.sample(&mut rng);
data[j].push(xsi);
}
}
let ef_construct = 25;
let nb_connection = 10;
let hnsw = Hnsw::<f32, DistL1>::new(nb_connection, nbcolumn, 16, ef_construct, DistL1 {});
for (i, d) in data.iter().enumerate() {
hnsw.insert((d, i));
}
hnsw.dump_layer_info();
let neighborhood_before_dump = FlatNeighborhood::from(&hnsw);
let nbg_2_before = neighborhood_before_dump.get_neighbours(2).unwrap();
println!("voisins du point 2 {:?}", nbg_2_before);
let fname = "dumpreloadtestflat";
let directory = tempfile::tempdir().unwrap();
let _res = hnsw.file_dump(directory.path(), fname);
debug!("HNSW reload");
let mut reloader = HnswIo::new(directory.path(), fname);
let hnsw_loaded: Hnsw<NoData, NoDist> = reloader.load_hnsw().unwrap();
let neighborhood_after_dump = FlatNeighborhood::from(&hnsw_loaded);
let nbg_2_after = neighborhood_after_dump.get_neighbours(2).unwrap();
println!("Neighbors of point 2 {:?}", nbg_2_after);
assert_eq!(nbg_2_after.len(), nbg_2_before.len());
for i in 0..nbg_2_before.len() {
assert_eq!(nbg_2_before[i].p_id, nbg_2_after[i].p_id);
assert_eq!(nbg_2_before[i].distance, nbg_2_after[i].distance);
}
check_graph_equality(&hnsw_loaded, &hnsw);
} }