Re-use deleted node IDs in incremental mode

src/metadata.rs

@@ -6,8 +6,6 @@ use ::roaring::RoaringBitmap;
 use byteorder::{BigEndian, ByteOrder};
 use heed::BoxedError;
 
-pub use crate::distance::Distance;
-pub use crate::error::Error;
 use crate::node::ItemIds;
 
 #[derive(Debug)]

src/parallel.rs

@@ -4,6 +4,7 @@ use std::io::{BufWriter, Write};
 use std::marker;
 use std::path::Path;
 use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use std::sync::Arc;
 
 use heed::types::Bytes;
 use heed::{BytesDecode, BytesEncode, RoTxn};
@@ -132,34 +133,100 @@ impl TmpNodesReader {
 /// A concurrent ID generate that will never return the same ID twice.
 #[derive(Debug)]
 pub struct ConcurrentNodeIds {
-    /// The current tree node ID we should use.
-    current: AtomicU32,
+    /// The current tree node ID we should use if there is no other IDs available.
+    current: Arc<AtomicU32>,
     /// The total number of tree node IDs used.
-    used: AtomicU64,
+    used: Arc<AtomicU64>,
+    /// A list of IDs to exhaust first before picking IDs from `current`.
+    /// When cloning this type, the available node ids are not cloned.
+    available: RoaringBitmap,
 }
 
 impl ConcurrentNodeIds {
-    /// Creates the ID generator starting at the given number.
+    /// Creates an ID generator returning unique IDs, avoiding the specified used IDs.
     pub fn new(used: RoaringBitmap) -> ConcurrentNodeIds {
         let last_id = used.iter().last().map(|id| id + 1).unwrap_or(0);
-        let used = used.len();
+        let used_ids = used.len();
+        let available = RoaringBitmap::from_sorted_iter(0..last_id).unwrap() - used;
 
-        ConcurrentNodeIds { current: AtomicU32::new(last_id), used: AtomicU64::new(used) }
+        ConcurrentNodeIds {
+            current: Arc::new(AtomicU32::new(last_id)),
+            used: Arc::new(AtomicU64::new(used_ids)),
+            available,
+        }
     }
 
-    /// Returns and increment the ID you can use as a NodeId.
-    pub fn next(&self) -> Result<u32> {
+    /// Returns a new unique ID and increase the count of IDs used.
+    pub fn next(&mut self) -> Result<u32> {
         if self.used.fetch_add(1, Ordering::Relaxed) > u32::MAX as u64 {
             Err(Error::DatabaseFull)
         } else {
-            Ok(self.current.fetch_add(1, Ordering::Relaxed))
+            match self.available.select(0) {
+                Some(id) => {
+                    // remove_smallest has a faster access pattern than remove
+                    self.available.remove_smallest(1);
+                    Ok(id)
+                }
+                None => Ok(self.current.fetch_add(1, Ordering::Relaxed)),
+            }
         }
     }
 
     /// Returns the number of used ids in total.
     pub fn used(&self) -> u64 {
         self.used.load(Ordering::Relaxed)
     }
+
+    /// Split a concurrent node ids into `n` parts..
+    pub fn split_in(self, n: usize) -> Vec<ConcurrentNodeIds> {
+        if n <= 1 {
+            return vec![self];
+        }
+
+        let Self { current, used, available } = self;
+        let chunk_size = available.len() / n as u64;
+        let mut ret = Vec::new();
+
+        let mut iter = available.into_iter();
+        for _ in 0..(n - 1) {
+            let available = iter.by_ref().take(chunk_size as usize).collect();
+            ret.push(Self { current: current.clone(), used: used.clone(), available });
+        }
+        // the last element is going to contain everything remaining
+        ret.push(Self { current, used, available: iter.collect() });
+
+        ret
+    }
+
+    /// Recompose a vector of concurrent node ids into a single concurrent node ids.
+    /// Return `None` if called on an empty list.
+    pub fn merge(concurrents_node_ids: Vec<ConcurrentNodeIds>) -> Option<Self> {
+        concurrents_node_ids.into_iter().reduce(|left, right| Self {
+            current: left.current,
+            used: left.used,
+            available: left.available | right.available,
+        })
+    }
+
+    /// Clone the concurrent node IDs **with** the available node IDs.
+    /// The available IDs should NEVER be shared between two different threads otherwise, they may
+    /// use the same ID for different nodes.
+    pub fn clone_with_available(&self) -> Self {
+        Self {
+            current: self.current.clone(),
+            used: self.used.clone(),
+            available: self.available.clone(),
+        }
+    }
+
+    /// Clone the concurrent node IDs **without** the available node IDs.
+    pub fn clone_without_available(&self) -> Self {
+        Self {
+            current: self.current.clone(),
+            used: self.used.clone(),
+            available: RoaringBitmap::new(),
+        }
+    }
 }
 
 /// A struture used to keep a list of the leaf nodes in the tree.

src/tests/writer.rs

@@ -859,12 +859,42 @@ fn reuse_node_id() {
     Item 3: Leaf(Leaf { header: NodeHeaderEuclidean { bias: 0.0 }, vector: [3.0000, 0.0000] })
     Item 4: Leaf(Leaf { header: NodeHeaderEuclidean { bias: 0.0 }, vector: [4.0000, 0.0000] })
     Item 5: Leaf(Leaf { header: NodeHeaderEuclidean { bias: 0.0 }, vector: [5.0000, 0.0000] })
+    Tree 0: Descendants(Descendants { descendants: [2, 3] })
     Tree 2: Descendants(Descendants { descendants: [4, 5] })
-    Tree 3: SplitPlaneNormal(SplitPlaneNormal { left: Tree(6), right: Tree(2), normal: [0.0000, 0.0000] })
+    Tree 3: SplitPlaneNormal(SplitPlaneNormal { left: Tree(5), right: Tree(2), normal: [0.0000, 0.0000] })
     Tree 4: SplitPlaneNormal(SplitPlaneNormal { left: Item(0), right: Tree(3), normal: [1.0000, 0.0000] })
-    Tree 5: Descendants(Descendants { descendants: [2, 3] })
-    Tree 6: SplitPlaneNormal(SplitPlaneNormal { left: Tree(5), right: Item(1), normal: [0.0000, 0.0000] })
+    Tree 5: SplitPlaneNormal(SplitPlaneNormal { left: Tree(0), right: Item(1), normal: [0.0000, 0.0000] })
     Root: Metadata { dimensions: 2, items: RoaringBitmap<[0, 1, 2, 3, 4, 5]>, roots: [4], distance: "euclidean" }
     updated_item_ids: RoaringBitmap<[]>
     "###);
+
+    let mut wtxn = handle.env.write_txn().unwrap();
+    let writer = Writer::new(handle.database, 0, 2).unwrap();
+
+    // if we now build a new tree, the id 1 should be re-used
+    writer.build(&mut wtxn, &mut rng, Some(2)).unwrap();
+    wtxn.commit().unwrap();
+
+    insta::assert_display_snapshot!(handle, @r###"
+    ==================
+    Dumping index 0
+    Item 0: Leaf(Leaf { header: NodeHeaderEuclidean { bias: 0.0 }, vector: [0.0000, 0.0000] })
+    Item 1: Leaf(Leaf { header: NodeHeaderEuclidean { bias: 0.0 }, vector: [1.0000, 0.0000] })
+    Item 2: Leaf(Leaf { header: NodeHeaderEuclidean { bias: 0.0 }, vector: [2.0000, 0.0000] })
+    Item 3: Leaf(Leaf { header: NodeHeaderEuclidean { bias: 0.0 }, vector: [3.0000, 0.0000] })
+    Item 4: Leaf(Leaf { header: NodeHeaderEuclidean { bias: 0.0 }, vector: [4.0000, 0.0000] })
+    Item 5: Leaf(Leaf { header: NodeHeaderEuclidean { bias: 0.0 }, vector: [5.0000, 0.0000] })
+    Tree 0: Descendants(Descendants { descendants: [2, 3] })
+    Tree 1: Descendants(Descendants { descendants: [3, 4] })
+    Tree 2: Descendants(Descendants { descendants: [4, 5] })
+    Tree 3: SplitPlaneNormal(SplitPlaneNormal { left: Tree(5), right: Tree(2), normal: [0.0000, 0.0000] })
+    Tree 4: SplitPlaneNormal(SplitPlaneNormal { left: Item(0), right: Tree(3), normal: [1.0000, 0.0000] })
+    Tree 5: SplitPlaneNormal(SplitPlaneNormal { left: Tree(0), right: Item(1), normal: [0.0000, 0.0000] })
+    Tree 6: SplitPlaneNormal(SplitPlaneNormal { left: Item(1), right: Tree(1), normal: [0.0000, 0.0000] })
+    Tree 7: Descendants(Descendants { descendants: [2, 5] })
+    Tree 8: SplitPlaneNormal(SplitPlaneNormal { left: Tree(6), right: Tree(7), normal: [0.0000, 0.0000] })
+    Tree 9: SplitPlaneNormal(SplitPlaneNormal { left: Tree(8), right: Item(0), normal: [-1.0000, 0.0000] })
+    Root: Metadata { dimensions: 2, items: RoaringBitmap<[0, 1, 2, 3, 4, 5]>, roots: [4, 9], distance: "euclidean" }
+    updated_item_ids: RoaringBitmap<[]>
+    "###);
 }