use the bias as intended

src/distance/angular.rs

@@ -1,5 +1,3 @@
-use std::borrow::Cow;
-
 use bytemuck::{Pod, Zeroable};
 use rand::Rng;
 
@@ -67,15 +65,15 @@ impl Distance for Angular {
     fn create_split<'a, R: Rng>(
         children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorCodec>>> {
+    ) -> heed::Result<Leaf<'static, Self>> {
         let [node_p, node_q] = two_means(rng, children, true)?;
         let vector: Vec<f32> =
             node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
         let unaligned_vector = UnalignedVector::from_vec(vector);
         let mut normal = Leaf { header: NodeHeaderAngular { norm: 0.0 }, vector: unaligned_vector };
         Self::normalize(&mut normal);
 
-        Ok(normal.vector)
+        Ok(normal)
     }
 
     fn margin_no_header(

src/distance/binary_quantized_angular.rs

@@ -1,5 +1,3 @@
-use std::borrow::Cow;
-
 use bytemuck::{Pod, Zeroable};
 use rand::Rng;
 
@@ -72,7 +70,7 @@ impl Distance for BinaryQuantizedAngular {
     fn create_split<'a, R: Rng>(
         children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorCodec>>> {
+    ) -> heed::Result<Leaf<'static, Self>> {
         let [node_p, node_q] = two_means::<Self, Angular, R>(rng, children, true)?;
         let vector: Vec<f32> =
             node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
@@ -83,7 +81,7 @@ impl Distance for BinaryQuantizedAngular {
         };
         Self::normalize(&mut normal);
 
-        Ok(normal.vector)
+        Ok(normal)
     }
 
     fn margin_no_header(

src/distance/binary_quantized_euclidean.rs

@@ -1,9 +1,7 @@
-use std::borrow::Cow;
-
 use bytemuck::{Pod, Zeroable};
 use rand::Rng;
 
-use super::{two_means_binary_quantized as two_means, Euclidean};
+use super::{two_means_binary_quantized as two_means, Euclidean, NodeHeaderEuclidean};
 use crate::distance::Distance;
 use crate::node::Leaf;
 use crate::parallel::ImmutableSubsetLeafs;
@@ -59,17 +57,35 @@
     fn create_split<'a, R: Rng>(
         children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorCodec>>> {
+    ) -> heed::Result<Leaf<'static, Self>> {
         let [node_p, node_q] = two_means::<Self, Euclidean, R>(rng, children, false)?;
         let vector: Vec<f32> =
             node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
         let mut normal = Leaf {
-            header: NodeHeaderBinaryQuantizedEuclidean { bias: 0.0 },
+            header: NodeHeaderBinaryQuantizedEuclidean::zeroed(),
             vector: UnalignedVector::from_slice(&vector),
         };
         Self::normalize(&mut normal);
 
-        Ok(Cow::Owned(normal.vector.into_owned()))
+        normal.header.bias = normal
+            .vector
+            .iter()
+            .zip(
+                UnalignedVector::<BinaryQuantized>::from_slice(
+                    &node_p.vector.iter().collect::<Vec<_>>(),
+                )
+                .iter(),
+            )
+            .zip(
+                UnalignedVector::<BinaryQuantized>::from_slice(
+                    &node_q.vector.iter().collect::<Vec<_>>(),
+                )
+                .iter(),
+            )
+            .map(|((n, p), q)| -n * (p + q) / 2.0)
+            .sum();
+
+        Ok(normal.into_owned())
     }
 
     fn margin(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {

src/distance/binary_quantized_manhattan.rs

@@ -1,9 +1,9 @@
 use std::borrow::Cow;
 
 use bytemuck::{Pod, Zeroable};
 use rand::Rng;
 
-use super::{two_means_binary_quantized as two_means, Manhattan};
+use super::{two_means_binary_quantized as two_means, Manhattan, NodeHeaderManhattan};
 use crate::distance::Distance;
 use crate::node::Leaf;
 use crate::parallel::ImmutableSubsetLeafs;
@@ -63,17 +63,35 @@
     fn create_split<'a, R: Rng>(
         children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorCodec>>> {
+    ) -> heed::Result<Leaf<'static, Self>> {
         let [node_p, node_q] = two_means::<Self, Manhattan, R>(rng, children, false)?;
         let vector: Vec<f32> =
             node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
         let mut normal = Leaf {
-            header: NodeHeaderBinaryQuantizedManhattan { bias: 0.0 },
+            header: NodeHeaderBinaryQuantizedManhattan::zeroed(),
             vector: UnalignedVector::from_slice(&vector),
         };
         Self::normalize(&mut normal);
 
-        Ok(Cow::Owned(normal.vector.into_owned()))
+        normal.header.bias = normal
+            .vector
+            .iter()
+            .zip(
+                UnalignedVector::<BinaryQuantized>::from_slice(
+                    &node_p.vector.iter().collect::<Vec<_>>(),
+                )
+                .iter(),
+            )
+            .zip(
+                UnalignedVector::<BinaryQuantized>::from_slice(
+                    &node_q.vector.iter().collect::<Vec<_>>(),
+                )
+                .iter(),
+            )
+            .map(|((n, p), q)| -n * (p + q) / 2.0)
+            .sum();
+
+        Ok(normal.into_owned())
     }
 
     fn margin(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {

src/distance/dot_product.rs

@@ -1,4 +1,4 @@
 use std::borrow::Cow;
 
 use bytemuck::{Pod, Zeroable};
 use heed::{RwPrefix, RwTxn};
@@ -90,7 +90,7 @@
     fn create_split<'a, R: Rng>(
         children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorCodec>>> {
+    ) -> heed::Result<Leaf<'static, Self>> {
         let [node_p, node_q] = two_means(rng, children, true)?;
         let vector: Vec<f32> =
             node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
@@ -101,7 +101,7 @@
         normal.header.extra_dim = node_p.header.extra_dim - node_q.header.extra_dim;
         Self::normalize(&mut normal);
 
-        Ok(normal.vector)
+        Ok(normal)
     }
 
     fn margin(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {

src/distance/euclidean.rs

@@ -1,5 +1,3 @@
-use std::borrow::Cow;
-
 use bytemuck::{Pod, Zeroable};
 use rand::Rng;
 
@@ -50,7 +48,7 @@ impl Distance for Euclidean {
     fn create_split<'a, R: Rng>(
         children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorCodec>>> {
+    ) -> heed::Result<Leaf<'static, Self>> {
         let [node_p, node_q] = two_means(rng, children, false)?;
         let vector: Vec<_> =
             node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
@@ -68,7 +66,7 @@ impl Distance for Euclidean {
             .map(|((n, p), q)| -n * (p + q) / 2.0)
             .sum();
 
-        Ok(normal.vector)
+        Ok(normal)
     }
 
     fn margin(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {

src/distance/manhattan.rs

@@ -1,4 +1,4 @@
 use std::borrow::Cow;
 
 use bytemuck::{Pod, Zeroable};
 use rand::Rng;
@@ -53,7 +53,7 @@
     fn create_split<'a, R: Rng>(
         children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorCodec>>> {
+    ) -> heed::Result<Leaf<'static, Self>> {
         let [node_p, node_q] = two_means(rng, children, false)?;
         let vector: Vec<_> =
             node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
@@ -71,7 +71,7 @@
             .map(|((n, p), q)| -n * (p + q) / 2.0)
             .sum();
 
-        Ok(normal.vector)
+        Ok(normal)
     }
 
     fn margin(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {

src/distance/mod.rs

@@ -1,4 +1,4 @@
 use std::borrow::Cow;
 use std::fmt;
 
 pub use angular::{Angular, NodeHeaderAngular};
@@ -30,7 +30,7 @@
 mod euclidean;
 mod manhattan;
 
-fn new_leaf<D: Distance>(vec: Vec<f32>) -> Leaf<'static, D> {
+pub fn new_leaf<D: Distance>(vec: Vec<f32>) -> Leaf<'static, D> {
     let vector = UnalignedVector::from_vec(vec);
     Leaf { header: D::new_header(&vector), vector }
 }
@@ -97,7 +97,7 @@
     fn create_split<'a, R: Rng>(
         children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorCodec>>>;
+    ) -> heed::Result<Leaf<'static, Self>>;
 
     fn margin(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {
         Self::margin_no_header(&p.vector, &q.vector)
@@ -108,12 +108,8 @@
         q: &UnalignedVector<Self::VectorCodec>,
     ) -> f32;
 
-    fn side<R: Rng>(
-        normal_plane: &UnalignedVector<Self::VectorCodec>,
-        node: &Leaf<Self>,
-        rng: &mut R,
-    ) -> Side {
-        let dot = Self::margin_no_header(&node.vector, normal_plane);
+    fn side<R: Rng>(normal_plane: &Leaf<Self>, node: &Leaf<Self>, rng: &mut R) -> Side {
+        let dot = Self::margin(normal_plane, node);
         if dot > 0.0 {
             Side::Right
         } else if dot < 0.0 {

src/node.rs

@@ -115,7 +115,7 @@ impl fmt::Debug for ItemIds<'_> {
 pub struct SplitPlaneNormal<'a, D: Distance> {
     pub left: NodeId,
     pub right: NodeId,
-    pub normal: Cow<'a, UnalignedVector<D::VectorCodec>>,
+    pub normal: Leaf<'a, D>,
 }
 
 impl<D: Distance> fmt::Debug for SplitPlaneNormal<'_, D> {
@@ -153,7 +153,8 @@ impl<'a, D: Distance> BytesEncode<'a> for NodeCodec<D> {
                 bytes.push(SPLIT_PLANE_NORMAL_TAG);
                 bytes.extend_from_slice(&left.to_bytes());
                 bytes.extend_from_slice(&right.to_bytes());
-                bytes.extend_from_slice(normal.as_bytes());
+                bytes.extend_from_slice(bytes_of(&normal.header));
+                bytes.extend_from_slice(normal.vector.as_bytes());
             }
             Node::Descendants(Descendants { descendants }) => {
                 bytes.push(DESCENDANTS_TAG);
@@ -179,8 +180,11 @@ impl<'a, D: Distance> BytesDecode<'a> for NodeCodec<D> {
             [SPLIT_PLANE_NORMAL_TAG, bytes @ ..] => {
                 let (left, bytes) = NodeId::from_bytes(bytes);
                 let (right, bytes) = NodeId::from_bytes(bytes);
+                let (header_bytes, remaining) = bytes.split_at(size_of::<D::Header>());
+                let header = pod_read_unaligned(header_bytes);
+                let vector = UnalignedVector::<D::VectorCodec>::from_bytes(remaining)?;
                 Ok(Node::SplitPlaneNormal(SplitPlaneNormal {
-                    normal: UnalignedVector::<D::VectorCodec>::from_bytes(bytes)?,
+                    normal: Leaf { header, vector },
                     left,
                     right,
                 }))

src/reader.rs

@@ -102,7 +102,7 @@
                 Node::SplitPlaneNormal(SplitPlaneNormal { normal, left, right }) => {
                     let left = recursive_depth(rtxn, database, index, left)?;
                     let right = recursive_depth(rtxn, database, index, right)?;
-                    let is_zero_normal = normal.is_zero() as usize;
+                    let is_zero_normal = normal.vector.is_zero() as usize;
 
                     Ok(TreeStats {
                         depth: 1 + left.depth.max(right.depth),
@@ -258,7 +258,7 @@
                     }
                 }
                 Node::SplitPlaneNormal(SplitPlaneNormal { normal, left, right }) => {
-                    let margin = D::margin_no_header(&normal, &query_leaf.vector);
+                    let margin = D::margin(&normal, &query_leaf);
                     queue.push((OrderedFloat(D::pq_distance(dist, margin, Side::Left)), left));
                     queue.push((OrderedFloat(D::pq_distance(dist, margin, Side::Right)), right));
                 }
@@ -327,7 +327,7 @@
                         writeln!(writer, "\t\t{} [label=\"{}\"]", key.node.item, key.node.item,)?
                     }
                     Node::SplitPlaneNormal(SplitPlaneNormal { normal, left, right }) => {
                         if normal.is_zero() {
                             writeln!(writer, "\t\t{} [color=red]", key.node.item)?;
                         }
                         writeln!(

src/writer.rs

@@ -10,7 +10,7 @@
 use rayon::prelude::*;
 use roaring::RoaringBitmap;
 
-use crate::distance::Distance;
+use crate::distance::{new_leaf, Distance};
 use crate::internals::{KeyCodec, Side};
 use crate::item_iter::ItemIter;
 use crate::node::{Descendants, ItemIds, Leaf, SplitPlaneNormal};
@@ -563,7 +563,7 @@
                         let mut left_ids = RoaringBitmap::new();
                         let mut right_ids = RoaringBitmap::new();
 
-                        if normal.is_zero() {
+                        if normal.vector.is_zero() {
                             randomly_split_children(rng, to_insert, &mut left_ids, &mut right_ids);
                         } else {
                             for leaf in to_insert {
@@ -730,7 +730,7 @@
                 let mut children_left = RoaringBitmap::new();
                 let mut children_right = RoaringBitmap::new();
                 randomly_split_children(rng, item_indices, &mut children_left, &mut children_right);
-                UnalignedVector::reset(&mut normal);
+                UnalignedVector::reset(&mut normal.vector);
 
                 (children_left, children_right)
             } else {