Implement basic progress tracking

This commit is contained in:
Dirkjan Ochtman 2020-12-15 11:51:00 +01:00
parent ee46e7cd1a
commit 71fd8cc15a
2 changed files with 25 additions and 0 deletions

View File

@ -6,6 +6,7 @@ edition = "2018"
[dependencies] [dependencies]
ahash = "0.6.1" ahash = "0.6.1"
indicatif = { version = "0.15", optional = true }
ordered-float = "2.0" ordered-float = "2.0"
rand = { version = "0.7.3", features = ["small_rng"] } rand = { version = "0.7.3", features = ["small_rng"] }
rayon = "1.5" rayon = "1.5"

View File

@ -3,6 +3,8 @@ use std::hash::Hash;
use std::ops::Index; use std::ops::Index;
use ahash::AHashSet as HashSet; use ahash::AHashSet as HashSet;
#[cfg(feature = "indicatif")]
use indicatif::ProgressBar;
use ordered_float::OrderedFloat; use ordered_float::OrderedFloat;
use rand::rngs::SmallRng; use rand::rngs::SmallRng;
use rand::{RngCore, SeedableRng}; use rand::{RngCore, SeedableRng};
@ -28,6 +30,8 @@ where
fn new(points: &[P], builder: Builder) -> (Self, Vec<PointId>) { fn new(points: &[P], builder: Builder) -> (Self, Vec<PointId>) {
let ef_search = builder.ef_search.unwrap_or(100); let ef_search = builder.ef_search.unwrap_or(100);
let ef_construction = builder.ef_construction.unwrap_or(100); let ef_construction = builder.ef_construction.unwrap_or(100);
#[cfg(feature = "indicatif")]
let progress = builder.progress;
if points.is_empty() { if points.is_empty() {
return ( return (
@ -99,6 +103,13 @@ where
for (layer, range) in ranges { for (layer, range) in ranges {
let num = if layer.0 > 0 { M } else { M * 2 }; let num = if layer.0 > 0 { M } else { M * 2 };
for &(_, pid) in &nodes[range] { for &(_, pid) in &nodes[range] {
#[cfg(feature = "indicatif")]
if pid.0 % 10_000 == 0 {
if let Some(bar) = &progress {
bar.set_position(pid.0 as u64);
}
}
search.reset(); search.reset();
let point = &points[pid]; let point = &points[pid];
search.push(PointId(0), &points[pid], &points); search.push(PointId(0), &points[pid], &points);
@ -128,6 +139,11 @@ where
} }
} }
#[cfg(feature = "indicatif")]
if let Some(bar) = progress {
bar.finish();
}
( (
Self { Self {
ef_search, ef_search,
@ -258,6 +274,8 @@ impl Default for Search {
pub struct Builder { pub struct Builder {
ef_search: Option<usize>, ef_search: Option<usize>,
ef_construction: Option<usize>, ef_construction: Option<usize>,
#[cfg(feature = "indicatif")]
progress: Option<ProgressBar>,
} }
impl Builder { impl Builder {
@ -274,6 +292,12 @@ impl Builder {
self self
} }
#[cfg(feature = "indicatif")]
pub fn progress(mut self, bar: ProgressBar) -> Self {
self.progress = Some(bar);
self
}
pub fn build<P: Point>(self, points: &[P]) -> (Hnsw<P>, Vec<PointId>) { pub fn build<P: Point>(self, points: &[P]) -> (Hnsw<P>, Vec<PointId>) {
Hnsw::new(points, self) Hnsw::new(points, self)
} }