chore: refactor to add proper DistanceWithTieBreak struct.

cevian · cevian · commit 460de910afeb · 2024-11-12T14:58:28.000-05:00
Previously this logic was all over the place. Now consolidated.
diff --git a/pgvectorscale/src/access_method/build.rs b/pgvectorscale/src/access_method/build.rs
@@ -864,11 +864,7 @@ pub mod tests {
         Ok(())
     }
 
-    pub fn verify_index_accuracy(
-        index_options: &str,
-        expected_cnt: i64,
-        dimensions: usize,
-    ) -> spi::Result<()> {
+    pub fn verify_index_accuracy(expected_cnt: i64, dimensions: usize) -> spi::Result<()> {
         let test_vec: Option<Vec<f32>> = Spi::get_one(&format!(
             "SELECT('{{' || array_to_string(array_agg(1.0), ',', '0') || '}}')::real[] AS embedding
     FROM generate_series(1, {dimensions})"
@@ -966,7 +962,8 @@ pub mod tests {
                         ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding
             FROM generate_series(1, {dimensions}));"))?;
 
-        verify_index_accuracy(index_options, expected_cnt, dimensions)
+        verify_index_accuracy(expected_cnt, dimensions)?;
+        Ok(())
     }
 
     #[pg_test]
@@ -1012,7 +1009,7 @@ pub mod tests {
                         ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding
             FROM generate_series(1, {dimensions}));"))?;
 
-        verify_index_accuracy(index_options, expected_cnt, dimensions)?;
+        verify_index_accuracy(expected_cnt, dimensions)?;
         Ok(())
     }
 }
diff --git a/pgvectorscale/src/access_method/graph.rs b/pgvectorscale/src/access_method/graph.rs
@@ -10,15 +10,15 @@ use crate::util::{HeapPointer, IndexPointer, ItemPointer};
 
 use super::graph_neighbor_store::GraphNeighborStore;
 
+use super::neighbor_with_distance::{Distance, DistanceWithTieBreak};
 use super::pg_vector::PgVector;
 use super::stats::{GreedySearchStats, InsertStats, PruneNeighborStats, StatsNodeVisit};
 use super::storage::Storage;
 use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance};
 
 pub struct ListSearchNeighbor<PD> {
     pub index_pointer: IndexPointer,
-    distance: f32,
-    distance_tie_break: usize, /* only used if distance = 0. This ensures a consistent order of results when distance = 0 */
+    distance_with_tie_break: DistanceWithTieBreak,
     private_data: PD,
 }
 
@@ -38,23 +38,21 @@ impl<PD> Eq for ListSearchNeighbor<PD> {}
 
 impl<PD> Ord for ListSearchNeighbor<PD> {
     fn cmp(&self, other: &Self) -> Ordering {
-        if self.distance == 0.0 && other.distance == 0.0 {
-            /* this logic should be consistent with what's used during pruning */
-            return self.distance_tie_break.cmp(&other.distance_tie_break);
-        }
-        self.distance.total_cmp(&other.distance)
+        self.distance_with_tie_break
+            .cmp(&other.distance_with_tie_break)
     }
 }
 
 impl<PD> ListSearchNeighbor<PD> {
-    pub fn new(index_pointer: IndexPointer, distance: f32, private_data: PD) -> Self {
-        assert!(!distance.is_nan());
-        debug_assert!(distance >= 0.0);
+    pub fn new(
+        index_pointer: IndexPointer,
+        distance_with_tie_break: DistanceWithTieBreak,
+        private_data: PD,
+    ) -> Self {
         Self {
             index_pointer,
             private_data,
-            distance,
-            distance_tie_break: 0,
+            distance_with_tie_break,
         }
     }
 
@@ -116,16 +114,22 @@ impl<QDM, PD> ListSearchResult<QDM, PD> {
         self.inserted.insert(ip)
     }
 
-    /// Internal function
-    pub fn insert_neighbor(&mut self, mut n: ListSearchNeighbor<PD>) {
-        self.stats.record_candidate();
-        if n.distance == 0.0 {
-            /* record the tie break if distance is 0 */
-            if let Some(tie_break_item_pointer) = self.tie_break_item_pointer {
-                let d = tie_break_item_pointer.ip_distance(n.index_pointer);
-                n.distance_tie_break = d;
+    pub fn create_distance_with_tie_break(
+        &self,
+        d: Distance,
+        ip: ItemPointer,
+    ) -> DistanceWithTieBreak {
+        match self.tie_break_item_pointer {
+            None => DistanceWithTieBreak::with_query(d, ip),
+            Some(tie_break_item_pointer) => {
+                DistanceWithTieBreak::new(d, ip, tie_break_item_pointer)
             }
         }
+    }
+
+    /// Internal function
+    pub fn insert_neighbor(&mut self, n: ListSearchNeighbor<PD>) {
+        self.stats.record_candidate();
         self.candidates.push(Reverse(n));
     }
 
@@ -332,8 +336,7 @@ impl<'a> Graph<'a> {
                     let list_search_entry = &lsr.visited[list_search_entry_idx];
                     visited_nodes.insert(NeighborWithDistance::new(
                         list_search_entry.index_pointer,
-                        list_search_entry.distance,
-                        list_search_entry.distance_tie_break,
+                        list_search_entry.distance_with_tie_break.clone(),
                     ));
                 }
             }
@@ -349,7 +352,7 @@ impl<'a> Graph<'a> {
     /// if we save the factors or the distances and add incrementally. Not sure.
     pub fn prune_neighbors<S: Storage>(
         &self,
-        neighbors_of: ItemPointer,
+        _neighbors_of: ItemPointer,
         mut candidates: Vec<NeighborWithDistance>,
         storage: &S,
         stats: &mut PruneNeighborStats,
@@ -410,8 +413,9 @@ impl<'a> Graph<'a> {
                         dist_state
                             .get_distance(candidate_neighbor.get_index_pointer_to_neighbor(), stats)
                     };
-                    let mut distance_between_candidate_and_point =
-                        candidate_neighbor.get_distance();
+                    let mut distance_between_candidate_and_point = candidate_neighbor
+                        .get_distance_with_tie_break()
+                        .get_distance();
 
                     //We need both values to be positive.
                     //Otherwise, the case where distance_between_candidate_and_point > 0 and distance_between_candidate_and_existing_neighbor < 0 is totally wrong.
@@ -451,8 +455,8 @@ impl<'a> Graph<'a> {
 
                             Note: with sbq these equivalence relations are actually not uncommon */
                             let ip_distance_between_candidate_and_point = candidate_neighbor
-                                .get_index_pointer_to_neighbor()
-                                .ip_distance(neighbors_of);
+                                .get_distance_with_tie_break()
+                                .get_distance_tie_break();
 
                             let ip_distance_between_candidate_and_existing_neighbor =
                                 candidate_neighbor
@@ -527,8 +531,7 @@ impl<'a> Graph<'a> {
             let (_needed_prune, contains) = self.update_back_pointer(
                 neighbor.get_index_pointer_to_neighbor(),
                 index_pointer,
-                neighbor.get_distance(),
-                neighbor.get_distance_tie_break(),
+                neighbor.get_distance_with_tie_break(),
                 storage,
                 &mut stats.prune_neighbor_stats,
             );
@@ -545,12 +548,14 @@ impl<'a> Graph<'a> {
         &mut self,
         from: IndexPointer,
         to: IndexPointer,
-        distance: f32,
-        distance_tie_break: usize,
+        distance_with_tie_break: &DistanceWithTieBreak,
         storage: &S,
         prune_stats: &mut PruneNeighborStats,
     ) -> (bool, bool) {
-        let new = vec![NeighborWithDistance::new(to, distance, distance_tie_break)];
+        let new = vec![NeighborWithDistance::new(
+            to,
+            distance_with_tie_break.clone(),
+        )];
         let (pruned, n) = self.add_neighbors(storage, from, new.clone(), prune_stats);
         (pruned, n.contains(&new[0]))
     }
diff --git a/pgvectorscale/src/access_method/neighbor_with_distance.rs b/pgvectorscale/src/access_method/neighbor_with_distance.rs
@@ -1,39 +1,104 @@
-use std::cmp::Ordering;
+use std::{cell::OnceCell, cmp::Ordering};
 
 use crate::util::{IndexPointer, ItemPointer};
 
 //TODO is this right?
 pub type Distance = f32;
+
+/* implements a distance with a lazy tie break */
 #[derive(Clone, Debug)]
-pub struct NeighborWithDistance {
-    index_pointer: IndexPointer,
+pub struct DistanceWithTieBreak {
     distance: Distance,
-    distance_tie_break: usize,
+    from: IndexPointer,
+    to: IndexPointer,
+    _distance_tie_break: OnceCell<usize>,
 }
 
-impl NeighborWithDistance {
-    pub fn new(
-        neighbor_index_pointer: ItemPointer,
-        distance: Distance,
-        distance_tie_break: usize,
-    ) -> Self {
+impl DistanceWithTieBreak {
+    pub fn new(distance: Distance, from: IndexPointer, to: IndexPointer) -> Self {
         assert!(!distance.is_nan());
         assert!(distance >= 0.0);
+        DistanceWithTieBreak {
+            distance,
+            from,
+            to,
+            _distance_tie_break: OnceCell::new(),
+        }
+    }
+
+    pub fn with_query(distance: Distance, to: IndexPointer) -> Self {
+        //this is the distance from the query to a index node.
+        //make the distance_tie_break = 0
+        let distance_tie_break = OnceCell::new();
+        distance_tie_break.set(0).unwrap();
+        DistanceWithTieBreak {
+            distance,
+            from: to,
+            to,
+            _distance_tie_break: distance_tie_break,
+        }
+    }
+
+    pub fn get_distance_tie_break(&self) -> usize {
+        *self
+            ._distance_tie_break
+            .get_or_init(|| self.from.ip_distance(self.to))
+    }
+
+    pub fn get_distance(&self) -> Distance {
+        self.distance
+    }
+}
+
+impl PartialOrd for DistanceWithTieBreak {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        if self.distance == 0.0 && other.distance == 0.0 {
+            return self
+                .get_distance_tie_break()
+                .partial_cmp(&other.get_distance_tie_break());
+        }
+        self.distance.partial_cmp(&other.distance)
+    }
+}
+
+impl Ord for DistanceWithTieBreak {
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.partial_cmp(other).unwrap()
+    }
+}
+
+impl PartialEq for DistanceWithTieBreak {
+    fn eq(&self, other: &Self) -> bool {
+        if self.distance == 0.0 && other.distance == 0.0 {
+            return self.get_distance_tie_break() == other.get_distance_tie_break();
+        }
+        self.distance == other.distance
+    }
+}
+
+//promise that PartialEq is reflexive
+impl Eq for DistanceWithTieBreak {}
+
+#[derive(Clone, Debug)]
+pub struct NeighborWithDistance {
+    index_pointer: IndexPointer,
+    distance: DistanceWithTieBreak,
+}
+
+impl NeighborWithDistance {
+    pub fn new(neighbor_index_pointer: ItemPointer, distance: DistanceWithTieBreak) -> Self {
         Self {
             index_pointer: neighbor_index_pointer,
             distance,
-            distance_tie_break,
         }
     }
 
     pub fn get_index_pointer_to_neighbor(&self) -> ItemPointer {
         self.index_pointer
     }
-    pub fn get_distance(&self) -> Distance {
-        self.distance
-    }
-    pub fn get_distance_tie_break(&self) -> usize {
-        return self.distance_tie_break;
+
+    pub fn get_distance_with_tie_break(&self) -> &DistanceWithTieBreak {
+        &self.distance
     }
 }
 
@@ -45,10 +110,7 @@ impl PartialOrd for NeighborWithDistance {
 
 impl Ord for NeighborWithDistance {
     fn cmp(&self, other: &Self) -> Ordering {
-        if self.distance == 0.0 && other.distance == 0.0 {
-            return self.distance_tie_break.cmp(&other.distance_tie_break);
-        }
-        self.distance.total_cmp(&other.distance)
+        self.distance.cmp(&other.distance)
     }
 }
 
diff --git a/pgvectorscale/src/access_method/plain_storage.rs b/pgvectorscale/src/access_method/plain_storage.rs
@@ -2,6 +2,7 @@ use super::{
     distance::DistanceFn,
     graph::{ListSearchNeighbor, ListSearchResult},
     graph_neighbor_store::GraphNeighborStore,
+    neighbor_with_distance::DistanceWithTieBreak,
     pg_vector::PgVector,
     plain_node::{ArchivedNode, Node, ReadableNode},
     stats::{
@@ -256,8 +257,7 @@ impl<'a> Storage for PlainStorage<'a> {
             let dist = unsafe { dist_state.get_distance(n, stats) };
             result.push(NeighborWithDistance::new(
                 n,
-                dist,
-                n.ip_distance(neighbors_of),
+                DistanceWithTieBreak::new(dist, neighbors_of, n),
             ))
         }
     }
@@ -288,7 +288,7 @@ impl<'a> Storage for PlainStorage<'a> {
 
         ListSearchNeighbor::new(
             index_pointer,
-            distance,
+            lsr.create_distance_with_tie_break(distance, index_pointer),
             PlainStorageLsnPrivateData::new(index_pointer, node, gns),
         )
     }
@@ -322,7 +322,7 @@ impl<'a> Storage for PlainStorage<'a> {
             };
             let lsn = ListSearchNeighbor::new(
                 neighbor_index_pointer,
-                distance,
+                lsr.create_distance_with_tie_break(distance, neighbor_index_pointer),
                 PlainStorageLsnPrivateData::new(neighbor_index_pointer, node_neighbor, gns),
             );
 
diff --git a/pgvectorscale/src/access_method/sbq.rs b/pgvectorscale/src/access_method/sbq.rs
