Implement basic progress tracking
This commit is contained in:
parent
ee46e7cd1a
commit
71fd8cc15a
|
@ -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"
|
||||||
|
|
24
src/lib.rs
24
src/lib.rs
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue