Simplify construction preparation

This commit is contained in:
Dirkjan Ochtman 2021-05-20 20:51:47 +02:00
parent 7288c9230a
commit 4a4ee88907
1 changed files with 20 additions and 47 deletions

View File

@ -235,6 +235,7 @@ where
} }
sizes.push((num, num)); sizes.push((num, num));
sizes.reverse(); sizes.reverse();
let top = LayerId(sizes.len() - 1);
// Give all points a random layer and sort the list of nodes by descending order for // Give all points a random layer and sort the list of nodes by descending order for
// construction. This allows us to copy higher layers to lower layers as construction // construction. This allows us to copy higher layers to lower layers as construction
@ -246,34 +247,15 @@ where
.collect::<Vec<_>>(); .collect::<Vec<_>>();
shuffled.sort_unstable(); shuffled.sort_unstable();
let mut new_points = Vec::with_capacity(points.len());
let mut new_nodes = Vec::with_capacity(points.len());
let mut out = vec![INVALID; points.len()]; let mut out = vec![INVALID; points.len()];
for (_, idx) in shuffled { let points = shuffled
let pid = PointId(new_nodes.len() as u32); .into_iter()
let layer = sizes
.iter()
.enumerate() .enumerate()
.find_map(|(i, &size)| match (pid.0 as usize) < size.1 { .map(|(i, (_, idx))| {
true => Some(i), out[idx] = PointId(i as u32);
false => None, points[idx].clone()
}) })
.unwrap(); .collect::<Vec<_>>();
new_points.push(points[idx].clone());
new_nodes.push((LayerId(sizes.len() - layer - 1), pid));
out[idx] = pid;
}
let (points, nodes) = (new_points, new_nodes);
debug_assert_eq!(nodes.last().unwrap().0, LayerId(0));
debug_assert_eq!(nodes.first().unwrap().0, LayerId(sizes.len() - 1));
// The layer from the first node is our top layer, or the zero layer if we have no nodes.
let top = match nodes.first() {
Some((top, _)) => *top,
None => LayerId(0),
};
// Figure out how many nodes will go on each layer. This helps us allocate memory capacity // Figure out how many nodes will go on each layer. This helps us allocate memory capacity
// for each layer in advance, and also helps enable batch insertion of points. // for each layer in advance, and also helps enable batch insertion of points.
@ -313,17 +295,15 @@ where
bar.set_message(format!("Building index (layer {})", layer.0)); bar.set_message(format!("Building index (layer {})", layer.0));
} }
let inserter = |pid| state.insert(pid, layer, &layers);
let end = range.end; let end = range.end;
if layer == top { if layer == top {
nodes[range].into_iter().for_each(|(_, pid)| { range.into_iter().for_each(|i| inserter(PointId(i as u32)))
let node = state.zero[*pid].write();
state.insert(*pid, node, layer, &layers);
})
} else { } else {
nodes[range].into_par_iter().for_each(|(_, pid)| { range
let node = state.zero[*pid].write(); .into_par_iter()
state.insert(*pid, node, layer, &layers); .for_each(|i| inserter(PointId(i as u32)));
});
} }
// For layers above the zero layer, make a copy of the current state of the zero layer // For layers above the zero layer, make a copy of the current state of the zero layer
@ -431,21 +411,14 @@ struct Construction<'a, P: Point> {
impl<'a, P: Point> Construction<'a, P> { impl<'a, P: Point> Construction<'a, P> {
/// Insert new node in the zero layer /// Insert new node in the zero layer
/// ///
/// * `new`: the `PointId` for the new node /// * `new` is the `PointId` for the new node
/// * `insertion`: a `Search` for shrinking a neighbor set (only used with heuristic neighbor selection)
/// * `search`: the result for searching potential neighbors for the new node
/// * `layer` contains all the nodes at the current layer /// * `layer` contains all the nodes at the current layer
/// * `points` is a slice of all the points in the index /// * `layers` refers to the existing higher-level layers
/// ///
/// Creates the new node, initializing its `nearest` array and updates the nearest neighbors /// Creates the new node, initializing its `nearest` array and updates the nearest neighbors
/// for the new node's neighbors if necessary before appending the new node to the layer. /// for the new node's neighbors if necessary before appending the new node to the layer.
fn insert( fn insert(&self, new: PointId, layer: LayerId, layers: &[Vec<UpperNode>]) {
&self, let mut node = self.zero[new].write();
new: PointId,
mut node: parking_lot::RwLockWriteGuard<ZeroNode>,
layer: LayerId,
layers: &[Vec<UpperNode>],
) {
let (mut search, mut insertion) = self.pool.pop(); let (mut search, mut insertion) = self.pool.pop();
insertion.ef = self.ef_construction; insertion.ef = self.ef_construction;