Merge pull request #130 from meilisearch/remove-sync-point

Parallelize arroy again*

Author: irevoire

Cargo.toml

@@ -28,6 +28,8 @@ thiserror = "2.0.9"
 nohash = "0.2.0"
 page_size = "0.6.0"
 enum-iterator = "2.1.0"
+thread_local = "1.1.8"
+crossbeam = "0.8.4"
 
 [dev-dependencies]
 anyhow = "1.0.95"

src/distance/cosine.rs

@@ -20,7 +20,7 @@ pub enum Cosine {}
 #[repr(C)]
 #[derive(Pod, Zeroable, Clone, Copy)]
 pub struct NodeHeaderCosine {
-    norm: f32,
+    pub(crate) norm: f32,
 }
 impl fmt::Debug for NodeHeaderCosine {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

src/distance/mod.rs

@@ -117,6 +117,10 @@ pub trait Distance: Send + Sync + Sized + Clone + fmt::Debug + 'static {
     ) -> heed::Result<()> {
         Ok(())
     }
+
+    fn size_of_item(dimensions: usize) -> usize {
+        std::mem::size_of::<Self::Header>() + Self::VectorCodec::size_of_item(dimensions)
+    }
 }
 
 fn two_means<D: Distance, R: Rng>(

src/metadata.rs

@@ -23,7 +23,7 @@ impl<'a> heed::BytesEncode<'a> for MetadataCodec {
 
     fn bytes_encode(item: &'a Self::EItem) -> Result<Cow<'a, [u8]>, BoxedError> {
         let Metadata { dimensions, items, roots, distance } = item;
-        debug_assert!(!distance.as_bytes().iter().any(|&b| b == 0));
+        debug_assert!(!distance.as_bytes().contains(&0));
 
         let mut output = Vec::with_capacity(
             size_of::<u32>()

src/node.rs

@@ -19,6 +19,20 @@ pub enum Node<'a, D: Distance> {
     SplitPlaneNormal(SplitPlaneNormal<'a, D>),
 }
 
+impl<'a, D: Distance> Node<'a, D> {
+    pub fn into_owned(self) -> Node<'static, D> {
+        match self {
+            Node::Leaf(leaf) => Node::Leaf(leaf.into_owned()),
+            Node::Descendants(descendants) => Node::Descendants(Descendants {
+                descendants: Cow::Owned(descendants.descendants.into_owned()),
+            }),
+            Node::SplitPlaneNormal(split_plane_normal) => {
+                Node::SplitPlaneNormal(split_plane_normal.into_owned())
+            }
+        }
+    }
+}
+
 /// A node generic over the version of the database.
 /// Should only be used while reading from the database.
 #[derive(Clone, Debug)]
@@ -40,8 +54,15 @@ impl<'a, D: Distance> Node<'a, D> {
             None
         }
     }
-}
 
+    pub fn descendants(self) -> Option<Descendants<'a>> {
+        if let Node::Descendants(descendants) = self {
+            Some(descendants)
+        } else {
+            None
+        }
+    }
+}
 /// A leaf node which corresponds to the vector inputed
 /// by the user and the distance header.
 pub struct Leaf<'a, D: Distance> {
@@ -142,6 +163,19 @@ impl<D: Distance> fmt::Debug for SplitPlaneNormal<'_, D> {
     }
 }
 
+impl<D: Distance> SplitPlaneNormal<'_, D> {
+    pub fn into_owned(self) -> SplitPlaneNormal<'static, D> {
+        SplitPlaneNormal {
+            left: self.left,
+            right: self.right,
+            normal: self.normal.map(|normal| Leaf {
+                header: normal.header,
+                vector: Cow::Owned(normal.vector.into_owned()),
+            }),
+        }
+    }
+}
+
 impl<D: Distance> Clone for SplitPlaneNormal<'_, D> {
     fn clone(&self) -> Self {
         Self { left: self.left, right: self.right, normal: self.normal.clone() }