use serde::{Deserialize, Serialize};
use cpu_time::ProcessTime;
use std::time::SystemTime;
use std::cmp::Ordering;
use parking_lot::{Mutex, RwLock, RwLockReadGuard};
use rayon::prelude::*;
use std::sync::Arc;
use std::sync::mpsc::channel;
use std::any::type_name;
use hashbrown::HashMap;
#[allow(unused)]
use std::collections::HashSet;
use std::collections::binary_heap::BinaryHeap;
use log::trace;
use log::{debug, info};
pub use crate::filter::FilterT;
use anndists::dist::distances::Distance;
#[derive(Default, Clone, Copy, Serialize, Deserialize, Debug)]
pub struct NoData;
pub(crate) const NB_LAYER_MAX: u8 = 16;
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct PointId(pub u8, pub i32);
pub type DataId = usize;
pub type PointDistance<T> = Box<dyn Distance<T>>;
#[derive(Debug, Clone, Copy)]
pub struct PointIdWithOrder {
pub point_id: PointId,
pub dist_to_ref: f32,
}
impl PartialEq for PointIdWithOrder {
fn eq(&self, other: &PointIdWithOrder) -> bool {
self.dist_to_ref == other.dist_to_ref
} }
impl PartialOrd for PointIdWithOrder {
fn partial_cmp(&self, other: &PointIdWithOrder) -> Option<Ordering> {
self.dist_to_ref.partial_cmp(&other.dist_to_ref)
} }
impl<T: Send + Sync + Clone + Copy> From<&PointWithOrder<'_, T>> for PointIdWithOrder {
fn from(point: &PointWithOrder<T>) -> PointIdWithOrder {
PointIdWithOrder::new(point.point_ref.p_id, point.dist_to_ref)
}
}
impl PointIdWithOrder {
pub fn new(point_id: PointId, dist_to_ref: f32) -> Self {
PointIdWithOrder {
point_id,
dist_to_ref,
}
}
}
#[repr(C)]
#[derive(Debug, Copy, Clone, Default)]
pub struct Neighbour {
pub d_id: DataId,
pub distance: f32,
pub p_id: PointId,
}
impl Neighbour {
pub fn new(d_id: DataId, distance: f32, p_id: PointId) -> Neighbour {
Neighbour {
d_id,
distance,
p_id,
}
}
pub fn get_origin_id(&self) -> DataId {
self.d_id
}
pub fn get_distance(&self) -> f32 {
self.distance
}
}
#[derive(Debug, Clone)]
enum PointData<'b, T: Clone + Send + Sync + 'b> {
V(Vec<T>),
S(&'b [T]),
}
impl<'b, T: Clone + Send + Sync + 'b> PointData<'b, T> {
fn new_v(v: Vec<T>) -> Self {
PointData::V(v)
}
fn new_s(s: &'b [T]) -> Self {
PointData::S(s)
}
fn get_v(&self) -> &[T] {
match self {
PointData::V(v) => v.as_slice(),
PointData::S(s) => s,
}
} }
#[derive(Debug, Clone)]
#[allow(clippy::type_complexity)]
pub struct Point<'b, T: Clone + Send + Sync> {
data: PointData<'b, T>,
origin_id: DataId,
p_id: PointId,
pub(crate) neighbours: Arc<RwLock<Vec<Vec<Arc<PointWithOrder<'b, T>>>>>>,
}
impl<'b, T: Clone + Send + Sync> Point<'b, T> {
pub fn new(v: Vec<T>, origin_id: usize, p_id: PointId) -> Self {
let mut neighbours = Vec::with_capacity(NB_LAYER_MAX as usize);
for _ in 0..NB_LAYER_MAX {
neighbours.push(Vec::<Arc<PointWithOrder<T>>>::new());
}
Point {
data: PointData::new_v(v),
origin_id,
p_id,
neighbours: Arc::new(RwLock::new(neighbours)),
}
}
pub fn new_from_mmap(s: &'b [T], origin_id: usize, p_id: PointId) -> Self {
let mut neighbours = Vec::with_capacity(NB_LAYER_MAX as usize);
for _ in 0..NB_LAYER_MAX {
neighbours.push(Vec::<Arc<PointWithOrder<T>>>::new());
}
Point {
data: PointData::new_s(s),
origin_id,
p_id,
neighbours: Arc::new(RwLock::new(neighbours)),
}
}
pub fn get_v(&self) -> &[T] {
self.data.get_v()
}
pub fn get_point_id(&self) -> PointId {
self.p_id
}
pub fn get_origin_id(&self) -> usize {
self.origin_id
}
pub fn get_neighborhood_id(&self) -> Vec<Vec<Neighbour>> {
let ref_neighbours = self.neighbours.read();
let nb_layer = ref_neighbours.len();
let mut neighborhood = Vec::<Vec<Neighbour>>::with_capacity(nb_layer);
for i in 0..nb_layer {
let mut neighbours = Vec::<Neighbour>::new();
let nb_ngbh = ref_neighbours[i].len();
if nb_ngbh > 0usize {
neighbours.reserve(nb_ngbh);
for pointwo in &ref_neighbours[i] {
neighbours.push(Neighbour::new(
pointwo.point_ref.get_origin_id(),
pointwo.dist_to_ref,
pointwo.point_ref.get_point_id(),
));
}
}
neighborhood.push(neighbours);
}
neighborhood
}
pub fn debug_dump(&self) {
println!(" \n dump of point id : {:?}", self.p_id);
println!("\n origin id : {:?} ", self.origin_id);
println!(" neighbours : ...");
let ref_neighbours = self.neighbours.read();
for i in 0..ref_neighbours.len() {
if !ref_neighbours[i].is_empty() {
println!("neighbours at layer {:?}", i);
for n in &ref_neighbours[i] {
println!(" {:?}", n.point_ref.p_id);
}
}
}
println!(" neighbours dump : end");
}
}
#[derive(Debug, Clone)]
pub(crate) struct PointWithOrder<'b, T: Clone + Send + Sync> {
point_ref: Arc<Point<'b, T>>,
dist_to_ref: f32,
}
impl<T: Clone + Send + Sync> PartialEq for PointWithOrder<'_, T> {
fn eq(&self, other: &PointWithOrder<T>) -> bool {
self.dist_to_ref == other.dist_to_ref
} }
impl<T: Clone + Send + Sync> Eq for PointWithOrder<'_, T> {}
#[allow(clippy::non_canonical_partial_ord_impl)]
impl<T: Clone + Send + Sync> PartialOrd for PointWithOrder<'_, T> {
fn partial_cmp(&self, other: &PointWithOrder<T>) -> Option<Ordering> {
self.dist_to_ref.partial_cmp(&other.dist_to_ref)
} }
impl<T: Clone + Send + Sync> Ord for PointWithOrder<'_, T> {
fn cmp(&self, other: &PointWithOrder<T>) -> Ordering {
if !self.dist_to_ref.is_nan() && !other.dist_to_ref.is_nan() {
self.dist_to_ref.partial_cmp(&other.dist_to_ref).unwrap()
} else {
panic!("got a NaN in a distance");
}
} }
impl<'b, T: Clone + Send + Sync> PointWithOrder<'b, T> {
pub fn new(point_ref: &Arc<Point<'b, T>>, dist_to_ref: f32) -> Self {
PointWithOrder {
point_ref: Arc::clone(point_ref),
dist_to_ref,
}
}
}
use rand::distr::Uniform;
use rand::prelude::*;
pub struct LayerGenerator {
rng: Arc<Mutex<rand::rngs::StdRng>>,
unif: Uniform<f64>,
scale: f64,
maxlevel: usize,
}
impl LayerGenerator {
pub fn new(max_nb_connection: usize, maxlevel: usize) -> Self {
let scale = 1. / (max_nb_connection as f64).ln();
LayerGenerator {
rng: Arc::new(Mutex::new(StdRng::from_os_rng())),
unif: Uniform::<f64>::new(0., 1.).unwrap(),
scale,
maxlevel,
}
}
pub(crate) fn new_with_scale(
max_nb_connection: usize,
scale_factor: f64,
maxlevel: usize,
) -> Self {
let scale_default = 1. / (max_nb_connection as f64).ln();
LayerGenerator {
rng: Arc::new(Mutex::new(StdRng::from_os_rng())),
unif: Uniform::<f64>::new(0., 1.).unwrap(),
scale: scale_default * scale_factor,
maxlevel,
}
}
fn generate(&self) -> usize {
let mut protected_rng = self.rng.lock();
let xsi = protected_rng.sample(self.unif);
let level = -xsi.ln() * self.scale;
let mut ulevel = level.floor() as usize;
if ulevel >= self.maxlevel {
ulevel = protected_rng.sample(Uniform::<usize>::new(0, self.maxlevel).unwrap());
}
ulevel
}
fn set_scale_modification(&mut self, scale_modification: f64) {
self.scale *= scale_modification;
log::info!("using scale for sampling levels : {:.2e}", self.scale);
}
fn get_level_scale(&self) -> f64 {
self.scale
}
}
type Layer<'b, T> = Vec<Arc<Point<'b, T>>>;
#[allow(unused)]
pub struct PointIndexation<'b, T: Clone + Send + Sync> {
pub(crate) max_nb_connection: usize,
pub(crate) max_layer: usize,
pub(crate) points_by_layer: Arc<RwLock<Vec<Layer<'b, T>>>>,
pub(crate) layer_g: LayerGenerator,
pub(crate) nb_point: Arc<RwLock<usize>>,
pub(crate) entry_point: Arc<RwLock<Option<Arc<Point<'b, T>>>>>,
}
impl<T: Clone + Send + Sync> Drop for PointIndexation<'_, T> {
fn drop(&mut self) {
let cpu_start = ProcessTime::now();
let sys_now = SystemTime::now();
info!("entering PointIndexation drop");
fn clear_neighborhoods<T: Clone + Send + Sync>(init: &Point<T>) {
let mut neighbours = init.neighbours.write();
let nb_layer = neighbours.len();
for l in 0..nb_layer {
neighbours[l].clear();
}
neighbours.clear();
}
if let Some(i) = self.entry_point.write().as_ref() {
clear_neighborhoods(i.as_ref());
}
let nb_level = self.get_max_level_observed();
for l in 0..=nb_level {
trace!("clearing layer {}", l);
let layer = &mut self.points_by_layer.write()[l as usize];
layer.into_par_iter().for_each(|p| clear_neighborhoods(p));
layer.clear();
}
debug!("clearing self.points_by_layer...");
drop(self.points_by_layer.write());
debug!("exiting PointIndexation drop");
info!(
" drop sys time(s) {:?} cpu time {:?}",
sys_now.elapsed().unwrap().as_secs(),
cpu_start.elapsed().as_secs()
);
} }
impl<'b, T: Clone + Send + Sync> PointIndexation<'b, T> {
pub fn new(max_nb_connection: usize, max_layer: usize, max_elements: usize) -> Self {
let mut points_by_layer = Vec::with_capacity(max_layer);
for i in 0..max_layer {
let s = 1. / (max_nb_connection as f64).ln();
let frac = (-(i as f64) / s).exp() - (-((i + 1) as f64) / s);
let expected_size = ((frac * max_elements as f64).round()) as usize;
points_by_layer.push(Vec::with_capacity(expected_size));
}
let layer_g = LayerGenerator::new(max_nb_connection, max_layer);
PointIndexation {
max_nb_connection,
max_layer,
points_by_layer: Arc::new(RwLock::new(points_by_layer)),
layer_g,
nb_point: Arc::new(RwLock::new(0)),
entry_point: Arc::new(RwLock::new(None)),
}
}
pub fn get_max_level_observed(&self) -> u8 {
let opt = self.entry_point.read();
match opt.as_ref() {
Some(arc_point) => arc_point.p_id.0,
None => 0,
}
}
pub fn get_level_scale(&self) -> f64 {
self.layer_g.get_level_scale()
}
fn debug_dump(&self) {
println!(" debug dump of PointIndexation");
let max_level_observed = self.get_max_level_observed();
for l in 0..=max_level_observed as usize {
println!(
" layer {} : length : {} ",
l,
self.points_by_layer.read()[l].len()
);
}
println!(" debug dump of PointIndexation end");
}
fn generate_new_point(&self, data: &[T], origin_id: usize) -> (Arc<Point<'b, T>>, usize) {
let level = self.layer_g.generate();
let new_point;
{
let mut points_by_layer_ref = self.points_by_layer.write();
let mut p_id = PointId(level as u8, -1);
p_id.1 = points_by_layer_ref[p_id.0 as usize].len() as i32;
let point = Point::new(data.to_vec(), origin_id, p_id);
new_point = Arc::new(point);
trace!("definitive pushing of point {:?}", p_id);
points_by_layer_ref[p_id.0 as usize].push(Arc::clone(&new_point));
} let nb_point;
{
let mut lock_nb_point = self.nb_point.write();
*lock_nb_point += 1;
nb_point = *lock_nb_point;
if nb_point % 50000 == 0 {
println!(" setting number of points {:?} ", nb_point);
}
}
trace!(" setting number of points {:?} ", *self.nb_point);
(Arc::clone(&new_point), nb_point)
}
fn check_entry_point(&self, new_point: &Arc<Point<'b, T>>) {
trace!("trying to get a lock on entry point");
let mut entry_point_ref = self.entry_point.write();
match entry_point_ref.as_ref() {
Some(arc_point) => {
if new_point.p_id.0 > arc_point.p_id.0 {
debug!("Hnsw , inserting entry point {:?} ", new_point.p_id);
debug!(
"PointIndexation insert setting max level from {:?} to {:?}",
arc_point.p_id.0, new_point.p_id.0
);
*entry_point_ref = Some(Arc::clone(new_point));
}
}
None => {
trace!("initializing entry point");
debug!("Hnsw , inserting entry point {:?} ", new_point.p_id);
*entry_point_ref = Some(Arc::clone(new_point));
}
}
}
pub fn get_nb_point(&self) -> usize {
*self.nb_point.read()
}
pub fn get_layer_nb_point(&self, layer: usize) -> usize {
let nb_layer = self.points_by_layer.read().len();
if layer < nb_layer {
self.points_by_layer.read()[layer].len()
} else {
0
}
}
pub fn get_data_dimension(&self) -> usize {
let ep = self.entry_point.read();
match ep.as_ref() {
Some(point) => point.get_v().len(),
None => 0,
}
}
pub fn get_point_data(&self, p_id: &PointId) -> Option<Vec<T>> {
if p_id.1 < 0 {
return None;
}
let p: usize = std::convert::TryFrom::try_from(p_id.1).unwrap();
let l = p_id.0 as usize;
if p_id.0 <= self.get_max_level_observed() && p < self.get_layer_nb_point(l) {
Some(self.points_by_layer.read()[l][p].get_v().to_vec())
} else {
None
}
}
#[allow(unused)]
pub(crate) fn get_point(&self, p_id: &PointId) -> Option<Arc<Point<'b, T>>> {
if p_id.1 < 0 {
return None;
}
let p: usize = std::convert::TryFrom::try_from(p_id.1).unwrap();
let l = p_id.0 as usize;
if p_id.0 <= self.get_max_level_observed() && p < self.get_layer_nb_point(l) {
Some(self.points_by_layer.read()[l][p].clone())
} else {
None
}
}
pub fn get_layer_iterator<'a>(&'a self, layer: usize) -> IterPointLayer<'a, 'b, T> {
IterPointLayer::new(self, layer)
} }
pub struct IterPoint<'a, 'b, T: Clone + Send + Sync + 'b> {
point_indexation: &'a PointIndexation<'b, T>,
pi_guard: RwLockReadGuard<'a, Vec<Layer<'b, T>>>,
layer: i64,
slot_in_layer: i64,
}
impl<'a, 'b, T: Clone + Send + Sync> IterPoint<'a, 'b, T> {
pub fn new(point_indexation: &'a PointIndexation<'b, T>) -> Self {
let pi_guard: RwLockReadGuard<Vec<Layer<'b, T>>> = point_indexation.points_by_layer.read();
IterPoint {
point_indexation,
pi_guard,
layer: -1,
slot_in_layer: -1,
}
}
}
impl<'b, T: Clone + Send + Sync> Iterator for IterPoint<'_, 'b, T> {
type Item = Arc<Point<'b, T>>;
fn next(&mut self) -> Option<Self::Item> {
if self.layer == -1 {
self.layer = 0;
self.slot_in_layer = 0;
}
if (self.slot_in_layer as usize) < self.pi_guard[self.layer as usize].len() {
let slot = self.slot_in_layer as usize;
self.slot_in_layer += 1;
Some(self.pi_guard[self.layer as usize][slot].clone())
} else {
self.slot_in_layer = 0;
self.layer += 1;
let entry_point_ref = self.point_indexation.entry_point.read();
let points_by_layer = self.point_indexation.points_by_layer.read();
let entry_point_level = entry_point_ref.as_ref().unwrap().p_id.0;
while (self.layer as u8) <= entry_point_level
&& points_by_layer[self.layer as usize].is_empty()
{
self.layer += 1;
}
if (self.layer as u8) <= entry_point_level {
let slot = self.slot_in_layer as usize;
self.slot_in_layer += 1;
Some(points_by_layer[self.layer as usize][slot].clone())
} else {
None
}
}
} }
impl<'a, 'b, T: Clone + Send + Sync> IntoIterator for &'a PointIndexation<'b, T> {
type Item = Arc<Point<'b, T>>;
type IntoIter = IterPoint<'a, 'b, T>;
fn into_iter(self) -> Self::IntoIter {
IterPoint::new(self)
}
}
pub struct IterPointLayer<'a, 'b, T: Clone + Send + Sync> {
_point_indexation: &'a PointIndexation<'b, T>,
pi_guard: RwLockReadGuard<'a, Vec<Layer<'b, T>>>,
layer: usize,
slot_in_layer: usize,
}
impl<'a, 'b, T: Clone + Send + Sync> IterPointLayer<'a, 'b, T> {
pub fn new(point_indexation: &'a PointIndexation<'b, T>, layer: usize) -> Self {
let pi_guard: RwLockReadGuard<Vec<Layer<'b, T>>> = point_indexation.points_by_layer.read();
IterPointLayer {
_point_indexation: point_indexation,
pi_guard,
layer,
slot_in_layer: 0,
}
}
}
impl<'b, T: Clone + Send + Sync + 'b> Iterator for IterPointLayer<'_, 'b, T> {
type Item = Arc<Point<'b, T>>;
fn next(&mut self) -> Option<Self::Item> {
if (self.slot_in_layer) < self.pi_guard[self.layer].len() {
let slot = self.slot_in_layer;
self.slot_in_layer += 1;
Some(self.pi_guard[self.layer][slot].clone())
} else {
None
}
} }
pub struct Hnsw<'b, T: Clone + Send + Sync + 'b, D: Distance<T>> {
pub(crate) ef_construction: usize,
pub(crate) max_nb_connection: usize,
pub(crate) extend_candidates: bool,
pub(crate) keep_pruned: bool,
pub(crate) max_layer: usize,
pub(crate) layer_indexed_points: PointIndexation<'b, T>,
#[allow(unused)]
pub(crate) data_dimension: usize,
pub(crate) dist_f: D,
pub(crate) searching: bool,
pub(crate) datamap_opt: bool,
}
impl<'b, T: Clone + Send + Sync, D: Distance<T> + Send + Sync> Hnsw<'b, T, D> {
pub fn new(
max_nb_connection: usize,
max_elements: usize,
max_layer: usize,
ef_construction: usize,
f: D,
) -> Self {
let adjusted_max_layer = (NB_LAYER_MAX as usize).min(max_layer);
let layer_indexed_points =
PointIndexation::<T>::new(max_nb_connection, adjusted_max_layer, max_elements);
let extend_candidates = false;
let keep_pruned = false;
if max_nb_connection > 256 {
println!("error max_nb_connection must be less equal than 256");
std::process::exit(1);
}
info!("Hnsw max_nb_connection {:?}", max_nb_connection);
info!("Hnsw nb elements {:?}", max_elements);
info!("Hnsw ef_construction {:?}", ef_construction);
info!("Hnsw distance {:?}", type_name::<D>());
info!("Hnsw extend candidates {:?}", extend_candidates);
Hnsw {
max_nb_connection,
ef_construction,
extend_candidates,
keep_pruned,
max_layer: adjusted_max_layer,
layer_indexed_points,
data_dimension: 0,
dist_f: f,
searching: false,
datamap_opt: false,
}
}
pub fn get_ef_construction(&self) -> usize {
self.ef_construction
}
pub fn get_max_level(&self) -> usize {
self.max_layer
}
pub fn get_max_level_observed(&self) -> u8 {
self.layer_indexed_points.get_max_level_observed()
}
pub fn get_max_nb_connection(&self) -> u8 {
self.max_nb_connection as u8
}
pub fn get_nb_point(&self) -> usize {
self.layer_indexed_points.get_nb_point()
}
pub fn set_searching_mode(&mut self, flag: bool) {
self.searching = flag;
}
pub fn get_distance_name(&self) -> String {
type_name::<D>().to_string()
}
pub fn set_keeping_pruned(&mut self, flag: bool) {
self.keep_pruned = flag;
}
pub fn get_distance(&self) -> &D {
&self.dist_f
}
pub fn set_extend_candidates(&mut self, flag: bool) {
self.extend_candidates = flag;
}
pub(crate) fn get_datamap_opt(&self) -> bool {
self.datamap_opt
}
pub fn modify_level_scale(&mut self, scale_modification: f64) {
if self.get_nb_point() > 0 {
println!(
"using modify_level_scale is possible at creation of a Hnsw structure to ensure coherence between runs"
)
}
let min_factor = 0.2;
println!(
"\n Current scale value : {:.2e}, Scale modification factor asked : {:.2e},(modification factor must be between {:.2e} and 1.)",
self.layer_indexed_points.layer_g.scale, scale_modification, min_factor
);
if scale_modification > 1. {
println!(
"\n Scale modification not applied, modification arg {:.2e} not valid , factor must be less than 1.)",
scale_modification
);
} else if scale_modification < min_factor {
println!(
"\n Scale modification arg {:.2e} not valid , factor must be greater than {:.2e}, using {:.2e})",
scale_modification, min_factor, min_factor
);
}
self.layer_indexed_points
.layer_g
.set_scale_modification(scale_modification.max(min_factor).min(1.));
}
fn search_layer(
&self,
point: &[T],
entry_point: Arc<Point<'b, T>>,
ef: usize,
layer: u8,
filter: Option<&dyn FilterT>,
) -> BinaryHeap<Arc<PointWithOrder<'b, T>>> {
trace!(
"entering search_layer with entry_point_id {:?} layer : {:?} ef {:?} ",
entry_point.p_id, layer, ef
);
let skiplist_size = ef.max(2);
let mut return_points = BinaryHeap::<Arc<PointWithOrder<T>>>::with_capacity(skiplist_size);
if self.layer_indexed_points.points_by_layer.read()[layer as usize].is_empty() {
trace!("search layer {:?}, empty layer", layer);
return return_points;
}
if entry_point.p_id.1 < 0 {
trace!("search layer negative point id : {:?}", entry_point.p_id);
return return_points;
}
let dist_to_entry_point = self.dist_f.eval(point, entry_point.data.get_v());
trace!(" distance to entry point: {:?} ", dist_to_entry_point);
let mut visited_point_id = HashMap::<PointId, Arc<Point<T>>>::new();
visited_point_id.insert(entry_point.p_id, Arc::clone(&entry_point));
let mut candidate_points =
BinaryHeap::<Arc<PointWithOrder<T>>>::with_capacity(skiplist_size);
candidate_points.push(Arc::new(PointWithOrder::new(
&entry_point,
-dist_to_entry_point,
)));
return_points.push(Arc::new(PointWithOrder::new(
&entry_point,
dist_to_entry_point,
)));
while !candidate_points.is_empty() {
let c = candidate_points.pop().unwrap();
let f = return_points.peek().unwrap();
assert!(f.dist_to_ref >= 0.);
assert!(c.dist_to_ref <= 0.);
trace!(
"Comparaing c : {:?} f : {:?}",
-(c.dist_to_ref),
f.dist_to_ref
);
if -(c.dist_to_ref) > f.dist_to_ref {
trace!(
"Fast return from search_layer, nb points : {:?} \n \t c {:?} \n \t f {:?} dists: {:?} {:?}",
return_points.len(),
c.point_ref.p_id,
f.point_ref.p_id,
-(c.dist_to_ref),
f.dist_to_ref
);
if filter.is_none() {
return return_points;
} else if return_points.len() >= ef {
return_points.retain(|p| {
filter
.as_ref()
.unwrap()
.hnsw_filter(&p.point_ref.get_origin_id())
});
}
}
let neighbours_c_l = &c.point_ref.neighbours.read()[layer as usize];
let c_pid = c.point_ref.p_id;
trace!(
" search_layer, {:?} has nb neighbours : {:?} ",
c_pid,
neighbours_c_l.len()
);
for e in neighbours_c_l {
if !visited_point_id.contains_key(&e.point_ref.p_id) {
visited_point_id.insert(e.point_ref.p_id, Arc::clone(&e.point_ref));
trace!(" visited insertion {:?}", e.point_ref.p_id);
let f_opt = return_points.peek();
if f_opt.is_none() {
debug!("return points empty when inserting {:?}", e.point_ref.p_id);
return return_points;
}
let f = f_opt.unwrap();
let e_dist_to_p = self.dist_f.eval(point, e.point_ref.data.get_v());
let f_dist_to_p = f.dist_to_ref;
if e_dist_to_p < f_dist_to_p || return_points.len() < ef {
let e_prime = Arc::new(PointWithOrder::new(&e.point_ref, e_dist_to_p));
trace!(
" inserting new candidate {:?}",
e_prime.point_ref.p_id
);
candidate_points
.push(Arc::new(PointWithOrder::new(&e.point_ref, -e_dist_to_p)));
if filter.is_none() {
return_points.push(Arc::clone(&e_prime));
} else {
let id: &usize = &e_prime.point_ref.get_origin_id();
if filter.as_ref().unwrap().hnsw_filter(id) {
if return_points.len() == 1 {
let only_id = return_points.peek().unwrap().point_ref.origin_id;
if !filter.as_ref().unwrap().hnsw_filter(&only_id) {
return_points.clear()
}
}
return_points.push(Arc::clone(&e_prime))
}
}
if return_points.len() > ef {
return_points.pop();
}
} }
} } trace!(
"return from search_layer, nb points : {:?}",
return_points.len()
);
return_points
}
#[inline]
pub fn insert(&self, datav_with_id: (&[T], usize)) {
self.insert_slice((datav_with_id.0, datav_with_id.1))
}
pub fn insert_slice(&self, data_with_id: (&[T], usize)) {
let (data, origin_id) = data_with_id;
let keep_pruned = self.keep_pruned;
let (new_point, point_rank) = self
.layer_indexed_points
.generate_new_point(data, origin_id);
trace!("Hnsw insert generated new point {:?} ", new_point.p_id);
let level = new_point.p_id.0;
let mut enter_point_copy = None;
let mut max_level_observed = 0;
{
if let Some(arc_point) = self.layer_indexed_points.entry_point.read().as_ref() {
enter_point_copy = Some(Arc::clone(arc_point));
if point_rank == 1 {
debug!(
"Hnsw stored first point , direct return {:?} ",
new_point.p_id
);
return;
}
max_level_observed = enter_point_copy.as_ref().unwrap().p_id.0;
}
}
if enter_point_copy.is_none() {
self.layer_indexed_points.check_entry_point(&new_point);
return;
}
let mut dist_to_entry = self
.dist_f
.eval(data, enter_point_copy.as_ref().unwrap().data.get_v());
for l in ((level + 1)..(max_level_observed + 1)).rev() {
let mut sorted_points = self.search_layer(
data,
Arc::clone(enter_point_copy.as_ref().unwrap()),
1,
l,
None,
);
trace!(
"in insert :search_layer layer {:?}, returned {:?} points ",
l,
sorted_points.len()
);
if sorted_points.len() > 1 {
panic!(
"in insert : search_layer layer {:?}, returned {:?} points ",
l,
sorted_points.len()
);
}
if let Some(ep) = sorted_points.pop() {
if new_point.neighbours.read()[l as usize].len()
< self.get_max_nb_connection() as usize
{
new_point.neighbours.write()[l as usize].push(Arc::clone(&ep));
}
let tmp_dist = self.dist_f.eval(data, ep.point_ref.data.get_v());
if tmp_dist < dist_to_entry {
enter_point_copy = Some(Arc::clone(&ep.point_ref));
dist_to_entry = tmp_dist;
}
} else {
trace!("layer still empty {} : got null list", l);
}
}
for l in (0..level + 1).rev() {
let ef = self.ef_construction;
let mut sorted_points = self.search_layer(
data,
Arc::clone(enter_point_copy.as_ref().unwrap()),
ef,
l,
None,
);
trace!(
"in insert :search_layer layer {:?}, returned {:?} points ",
l,
sorted_points.len()
);
sorted_points = from_positive_binaryheap_to_negative_binary_heap(&mut sorted_points);
if !sorted_points.is_empty() {
let nb_conn;
let extend_c;
if l == 0 {
nb_conn = 2 * self.max_nb_connection;
extend_c = self.extend_candidates;
} else {
nb_conn = self.max_nb_connection;
extend_c = false;
}
let mut neighbours = Vec::<Arc<PointWithOrder<T>>>::with_capacity(nb_conn);
self.select_neighbours(
data,
&mut sorted_points,
nb_conn,
extend_c,
l,
keep_pruned,
&mut neighbours,
);
neighbours.sort_unstable();
new_point.neighbours.write()[l as usize].clone_from(&neighbours);
if !neighbours.is_empty() {
enter_point_copy = Some(Arc::clone(&neighbours[0].point_ref));
}
}
} self.reverse_update_neighborhood_simple(Arc::clone(&new_point));
self.layer_indexed_points.check_entry_point(&new_point);
trace!("Hnsw exiting insert new point {:?} ", new_point.p_id);
}
pub fn parallel_insert(&self, datas: &[(&Vec<T>, usize)]) {
debug!("entering parallel_insert");
datas
.par_iter()
.for_each(|&(item, v)| self.insert((item.as_slice(), v)));
debug!("exiting parallel_insert");
}
pub fn parallel_insert_slice(&self, datas: &Vec<(&[T], usize)>) {
datas.par_iter().for_each(|&item| self.insert_slice(item));
}
fn reverse_update_neighborhood_simple(&self, new_point: Arc<Point<T>>) {
trace!(
"reverse update neighbourhood for new point {:?} ",
new_point.p_id
);
let level = new_point.p_id.0;
for l in (0..level + 1).rev() {
for q in &new_point.neighbours.read()[l as usize] {
if new_point.p_id != q.point_ref.p_id {
let q_point = &q.point_ref;
let mut q_point_neighbours = q_point.neighbours.write();
let n_to_add = PointWithOrder::<T>::new(&Arc::clone(&new_point), q.dist_to_ref);
let l_n = n_to_add.point_ref.p_id.0 as usize;
let already = q_point_neighbours[l_n]
.iter()
.position(|old| old.point_ref.p_id == new_point.p_id);
if already.is_some() {
continue;
}
q_point_neighbours[l_n].push(Arc::new(n_to_add));
let nbn_at_l = q_point_neighbours[l_n].len();
let threshold_shrinking = if l_n > 0 {
self.max_nb_connection
} else {
2 * self.max_nb_connection
};
let shrink = nbn_at_l > threshold_shrinking;
{
q_point_neighbours[l_n].sort_unstable();
if shrink {
q_point_neighbours[l_n].pop();
}
}
} }
}
}
pub fn get_point_indexation(&self) -> &PointIndexation<'b, T> {
&self.layer_indexed_points
}
#[allow(clippy::too_many_arguments)]
fn select_neighbours(
&self,
data: &[T],
candidates: &mut BinaryHeap<Arc<PointWithOrder<'b, T>>>,
nb_neighbours_asked: usize,
extend_candidates_asked: bool,
layer: u8,
keep_pruned: bool,
neighbours_vec: &mut Vec<Arc<PointWithOrder<'b, T>>>,
) {
trace!(
"entering select_neighbours : nb candidates: {}",
candidates.len()
);
neighbours_vec.clear();
let mut extend_candidates = false;
if candidates.len() <= nb_neighbours_asked {
if !extend_candidates_asked {
while !candidates.is_empty() {
let p = candidates.pop().unwrap();
assert!(-p.dist_to_ref >= 0.);
neighbours_vec
.push(Arc::new(PointWithOrder::new(&p.point_ref, -p.dist_to_ref)));
}
return;
} else {
extend_candidates = true;
}
}
if extend_candidates {
let mut candidates_set = HashMap::<PointId, Arc<Point<T>>>::new();
for c in candidates.iter() {
candidates_set.insert(c.point_ref.p_id, Arc::clone(&c.point_ref));
}
let mut new_candidates_set = HashMap::<PointId, Arc<Point<T>>>::new();
for (_p_id, p_point) in candidates_set.iter() {
let n_p_layer = &p_point.neighbours.read()[layer as usize];
for q in n_p_layer {
if !candidates_set.contains_key(&q.point_ref.p_id)
&& !new_candidates_set.contains_key(&q.point_ref.p_id)
{
new_candidates_set.insert(q.point_ref.p_id, Arc::clone(&q.point_ref));
}
}
} trace!(
"select neighbours extend candidates from : {:?} adding : {:?}",
candidates.len(),
new_candidates_set.len()
);
for (_p_id, p_point) in new_candidates_set.iter() {
let dist_topoint = self.dist_f.eval(data, p_point.data.get_v());
candidates.push(Arc::new(PointWithOrder::new(p_point, -dist_topoint)));
}
} let mut discarded_points = BinaryHeap::<Arc<PointWithOrder<T>>>::new();
while !candidates.is_empty() && neighbours_vec.len() < nb_neighbours_asked {
if let Some(e_p) = candidates.pop() {
let mut e_to_insert = true;
let e_point_v = e_p.point_ref.data.get_v();
assert!(e_p.dist_to_ref <= 0.);
if !neighbours_vec.is_empty() {
e_to_insert = !neighbours_vec.iter().any(|d| {
self.dist_f.eval(e_point_v, d.point_ref.data.get_v()) <= -e_p.dist_to_ref
});
}
if e_to_insert {
trace!("inserting neighbours : {:?} ", e_p.point_ref.p_id);
neighbours_vec.push(Arc::new(PointWithOrder::new(
&e_p.point_ref,
-e_p.dist_to_ref,
)));
} else {
trace!("discarded neighbours : {:?} ", e_p.point_ref.p_id);
if keep_pruned {
discarded_points.push(Arc::new(PointWithOrder::new(
&e_p.point_ref,
e_p.dist_to_ref,
)));
}
}
}
}
if keep_pruned {
while !discarded_points.is_empty() && neighbours_vec.len() < nb_neighbours_asked {
let best_point = discarded_points.pop().unwrap();
assert!(best_point.dist_to_ref <= 0.);
neighbours_vec.push(Arc::new(PointWithOrder::new(
&best_point.point_ref,
-best_point.dist_to_ref,
)));
}
};
if log::log_enabled!(log::Level::Trace) {
trace!(
"exiting select_neighbours : nb candidates: {}",
neighbours_vec.len()
);
for n in neighbours_vec {
trace!(" neighbours {:?} ", n.point_ref.p_id);
}
}
}
pub fn dump_layer_info(&self) {
self.layer_indexed_points.debug_dump();
}
#[allow(unused)]
fn search_general(&self, data: &[T], knbn: usize, ef_arg: usize) -> Vec<Neighbour> {
let mut entry_point;
{
let entry_point_opt_ref = self.layer_indexed_points.entry_point.read();
if entry_point_opt_ref.is_none() {
return Vec::<Neighbour>::new();
} else {
entry_point = Arc::clone((*entry_point_opt_ref).as_ref().unwrap());
}
}
let mut dist_to_entry = self.dist_f.eval(data, entry_point.as_ref().data.get_v());
for layer in (1..=entry_point.p_id.0).rev() {
let mut neighbours = self.search_layer(data, Arc::clone(&entry_point), 1, layer, None);
neighbours = from_positive_binaryheap_to_negative_binary_heap(&mut neighbours);
if let Some(entry_point_tmp) = neighbours.pop() {
let tmp_dist = self
.dist_f
.eval(data, entry_point_tmp.point_ref.data.get_v());
if tmp_dist < dist_to_entry {
entry_point = Arc::clone(&entry_point_tmp.point_ref);
dist_to_entry = tmp_dist;
}
}
}
let ef = ef_arg.max(knbn);
let neighbours_heap = self.search_layer(data, entry_point, ef, 0, None);
let neighbours = neighbours_heap.into_sorted_vec();
let last = knbn.min(ef).min(neighbours.len());
let knn_neighbours: Vec<Neighbour> = neighbours[0..last]
.iter()
.map(|p| {
Neighbour::new(
p.as_ref().point_ref.origin_id,
p.as_ref().dist_to_ref,
p.as_ref().point_ref.p_id,
)
})
.collect();
knn_neighbours
}
pub fn search_filter(
&self,
data: &[T],
knbn: usize,
ef_arg: usize,
filter: Option<&dyn FilterT>,
) -> Vec<Neighbour> {
let entry_point;
{
let entry_point_opt_ref = self.layer_indexed_points.entry_point.read();
if entry_point_opt_ref.is_none() {
return Vec::<Neighbour>::new();
} else {
entry_point = Arc::clone((*entry_point_opt_ref).as_ref().unwrap());
}
}
let mut dist_to_entry = self.dist_f.eval(data, entry_point.as_ref().data.get_v());
let mut pivot = Arc::clone(&entry_point);
let mut new_pivot = None;
for layer in (1..=entry_point.p_id.0).rev() {
let mut has_changed = false;
{
let neighbours = &pivot.neighbours.read()[layer as usize];
for n in neighbours {
let tmp_dist = self.dist_f.eval(data, n.point_ref.data.get_v());
if tmp_dist < dist_to_entry {
new_pivot = Some(Arc::clone(&n.point_ref));
has_changed = true;
dist_to_entry = tmp_dist;
}
} }
if has_changed {
pivot = Arc::clone(new_pivot.as_ref().unwrap());
}
} let ef = ef_arg.max(knbn);
log::debug!("pivot changed , current pivot {:?}", pivot.get_point_id());
let mut l = 0u8;
let layer_to_search = loop {
if self.get_point_indexation().get_layer_nb_point(l as usize) > 0 {
break l;
}
l += 1;
};
let neighbours_heap = self.search_layer(data, pivot, ef, layer_to_search, filter);
let neighbours = neighbours_heap.into_sorted_vec();
let last = knbn.min(ef).min(neighbours.len());
if let Some(filter_t) = filter {
let knn_neighbours: Vec<Neighbour> = neighbours[0..last]
.iter()
.map(|p| {
if filter_t.hnsw_filter(&p.as_ref().point_ref.origin_id) {
Some(Neighbour::new(
p.as_ref().point_ref.origin_id,
p.as_ref().dist_to_ref,
p.as_ref().point_ref.p_id,
))
} else {
None
}
})
.filter(|x| x.is_some())
.map(|x| x.unwrap())
.collect();
knn_neighbours
} else {
let knn_neighbours: Vec<Neighbour> = neighbours[0..last]
.iter()
.map(|p| {
Neighbour::new(
p.as_ref().point_ref.origin_id,
p.as_ref().dist_to_ref,
p.as_ref().point_ref.p_id,
)
})
.collect();
knn_neighbours
}
}
#[inline]
pub fn search_possible_filter(
&self,
data: &[T],
knbn: usize,
ef_arg: usize,
filter: Option<&dyn FilterT>,
) -> Vec<Neighbour> {
self.search_filter(data, knbn, ef_arg, filter)
}
pub fn search(&self, data: &[T], knbn: usize, ef_arg: usize) -> Vec<Neighbour> {
self.search_possible_filter(data, knbn, ef_arg, None)
}
fn search_with_id(
&self,
request: (usize, &Vec<T>),
knbn: usize,
ef: usize,
) -> (usize, Vec<Neighbour>) {
(request.0, self.search(request.1, knbn, ef))
}
pub fn parallel_search(&self, datas: &[Vec<T>], knbn: usize, ef: usize) -> Vec<Vec<Neighbour>> {
let (sender, receiver) = channel();
let nb_request = datas.len();
let requests: Vec<(usize, &Vec<T>)> = (0..nb_request).zip(datas.iter()).collect();
requests.par_iter().for_each_with(sender, |s, item| {
s.send(self.search_with_id(*item, knbn, ef)).unwrap()
});
let req_res: Vec<(usize, Vec<Neighbour>)> = receiver.iter().collect();
let mut answers = Vec::<Vec<Neighbour>>::with_capacity(datas.len());
let mut req_hash = HashMap::<usize, usize>::new();
for (i, elt) in req_res.iter().enumerate() {
req_hash.insert(elt.0, i);
}
for i in 0..datas.len() {
let answer_i = req_hash.get_key_value(&i).unwrap().1;
answers.push((req_res[*answer_i].1).clone());
}
answers
} }
#[allow(unused)]
fn from_negative_binaryheap_to_sorted_vector<'b, T: Send + Sync + Copy>(
heap_points: &mut BinaryHeap<Arc<PointWithOrder<'b, T>>>,
) -> Vec<Arc<PointWithOrder<'b, T>>> {
let nb_points = heap_points.len();
let mut vec_points = Vec::<Arc<PointWithOrder<T>>>::with_capacity(nb_points);
for p in heap_points.iter() {
assert!(p.dist_to_ref <= 0.);
let reverse_p = Arc::new(PointWithOrder::new(&p.point_ref, -p.dist_to_ref));
vec_points.push(reverse_p);
}
trace!(
"from_negative_binaryheap_to_sorted_vector nb points in out {:?} {:?} ",
nb_points,
vec_points.len()
);
vec_points
}
fn from_positive_binaryheap_to_negative_binary_heap<'b, T: Send + Sync + Clone>(
positive_heap: &mut BinaryHeap<Arc<PointWithOrder<'b, T>>>,
) -> BinaryHeap<Arc<PointWithOrder<'b, T>>> {
let nb_points = positive_heap.len();
let mut negative_heap = BinaryHeap::<Arc<PointWithOrder<T>>>::with_capacity(nb_points);
for p in positive_heap.iter() {
assert!(p.dist_to_ref >= 0.);
let reverse_p = Arc::new(PointWithOrder::new(&p.point_ref, -p.dist_to_ref));
negative_heap.push(reverse_p);
}
trace!(
"from_positive_binaryheap_to_negative_binary_heap nb points in out {:?} {:?} ",
nb_points,
negative_heap.len()
);
negative_heap
}
#[allow(unused)]
pub(crate) fn check_graph_equality<T1, D1, T2, D2>(hnsw1: &Hnsw<T1, D1>, hnsw2: &Hnsw<T2, D2>)
where
T1: Copy + Clone + Send + Sync,
D1: Distance<T1> + Default + Send + Sync,
T2: Copy + Clone + Send + Sync,
D2: Distance<T2> + Default + Send + Sync,
{
debug!("In check_graph_equality");
assert_eq!(hnsw1.get_nb_point(), hnsw2.get_nb_point());
assert!(
hnsw1.layer_indexed_points.entry_point.read().is_some()
|| hnsw1.layer_indexed_points.entry_point.read().is_some(),
"one entry point is None"
);
let ep1_read = hnsw1.layer_indexed_points.entry_point.read();
let ep2_read = hnsw2.layer_indexed_points.entry_point.read();
let ep1 = ep1_read.as_ref().unwrap();
let ep2 = ep2_read.as_ref().unwrap();
assert_eq!(
ep1.origin_id, ep2.origin_id,
"different entry points {:?} {:?}",
ep1.origin_id, ep2.origin_id
);
assert_eq!(ep1.p_id, ep2.p_id, "origin id {:?} ", ep1.origin_id);
let layers_1 = hnsw1.layer_indexed_points.points_by_layer.read();
let layers_2 = hnsw2.layer_indexed_points.points_by_layer.read();
let mut nb_point_checked = 0;
let mut nb_neighbours_checked = 0;
for i in 0..NB_LAYER_MAX as usize {
debug!("Checking layer {:?}", i);
assert_eq!(layers_1[i].len(), layers_2[i].len());
for j in 0..layers_1[i].len() {
let p1 = &layers_1[i][j];
let p2 = &layers_2[i][j];
assert_eq!(p1.origin_id, p2.origin_id);
assert_eq!(
p1.p_id, p2.p_id,
"Checking origin_id point {:?} ",
p1.origin_id
);
nb_point_checked += 1;
let nbgh1 = p1.neighbours.read();
let nbgh2 = p2.neighbours.read();
assert_eq!(nbgh1.len(), nbgh2.len());
for k in 0..nbgh1.len() {
assert_eq!(nbgh1[k].len(), nbgh2[k].len());
for l in 0..nbgh1[k].len() {
assert_eq!(
nbgh1[k][l].point_ref.origin_id,
nbgh2[k][l].point_ref.origin_id
);
assert_eq!(nbgh1[k][l].point_ref.p_id, nbgh2[k][l].point_ref.p_id);
assert_eq!(nbgh1[k][l].dist_to_ref, nbgh2[k][l].dist_to_ref);
nb_neighbours_checked += 1;
}
}
} } assert_eq!(nb_point_checked, hnsw1.get_nb_point());
debug!("nb neighbours checked {:?}", nb_neighbours_checked);
debug!("exiting check_equality");
}
#[cfg(test)]
mod tests {
use super::*;
use anndists::dist;
fn log_init_test() {
let _ = env_logger::builder().is_test(true).try_init();
}
#[test]
fn test_iter_point() {
println!("\n\n test_iter_point");
let mut rng = rand::rng();
let unif = Uniform::<f32>::new(0., 1.).unwrap();
let nbcolumn = 5000;
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 = rng.sample(unif);
data[j].push(xsi);
}
}
let ef_construct = 25;
let nb_connection = 10;
let start = ProcessTime::now();
let hns = Hnsw::<f32, dist::DistL1>::new(
nb_connection,
nbcolumn,
16,
ef_construct,
dist::DistL1 {},
);
for (i, d) in data.iter().enumerate() {
hns.insert((d, i));
}
let cpu_time = start.elapsed();
println!(" test_insert_iter_point time inserting {:?}", cpu_time);
hns.dump_layer_info();
let ptiter = hns.get_point_indexation().into_iter();
let mut nb_dumped = 0;
for _point in ptiter {
nb_dumped += 1;
}
assert_eq!(nb_dumped, nbcolumn);
}
#[test]
fn test_iter_layerpoint() {
println!("\n\n test_iter_point");
let mut rng = rand::rng();
let unif = Uniform::<f32>::new(0., 1.).unwrap();
let nbcolumn = 5000;
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 = rng.sample(unif);
data[j].push(xsi);
}
}
let ef_construct = 25;
let nb_connection = 10;
let start = ProcessTime::now();
let hns = Hnsw::<f32, dist::DistL1>::new(
nb_connection,
nbcolumn,
16,
ef_construct,
dist::DistL1 {},
);
for (i, d) in data.iter().enumerate() {
hns.insert((d, i));
}
let cpu_time = start.elapsed();
println!(" test_insert_iter_point time inserting {:?}", cpu_time);
hns.dump_layer_info();
let layer_num = 0;
let nbpl = hns.get_point_indexation().get_layer_nb_point(layer_num);
let layer_iter = hns.get_point_indexation().get_layer_iterator(layer_num);
let mut nb_dumped = 0;
for _point in layer_iter {
nb_dumped += 1;
}
println!(
"test_iter_layerpoint : nb point in layer {} , nb found {}",
nbpl, nb_dumped
);
assert_eq!(nb_dumped, nbpl);
}
#[test]
fn test_sparse_search() {
log_init_test();
for _ in 0..800 {
let hnsw: Hnsw<f32, dist::DistL1> =
Hnsw::new(15, 100_000, 20, 500_000, dist::DistL1 {});
hnsw.insert((&[1.0, 0.0, 0.0, 0.0], 0));
let result = hnsw.search(&[1.0, 0.0, 0.0, 0.0], 2, 10);
assert_eq!(result, vec![Neighbour::new(0, 0.0, PointId(0, 0))]);
}
}
}