Fixes for poor recall with label filters (#225)

This PR fixes a pair of issues that were observed in benchmarks to cause
poor recall with label filtering in certain cases:
- missing check to preserve edges in pruning containing label
information not covered by other edges (cf. line 6 of Algorithm 3 from
Filtered DiskANN paper)
- missing logic to suppress label-based checks when inserting into the
default start node (resulting in poor recall when no label filter
condition was present)

Testing:
- new sanity check unit test
- validated benchmark runs to be reported in an upcoming blog post

---------

Signed-off-by: tjgreen42 <tj@timescale.com>
Co-authored-by: Matvey Arye <cevian@gmail.com>

Author: tjgreen42

pgvectorscale/src/access_method/build.rs

@@ -178,6 +178,10 @@ unsafe fn aminsert_internal(
     }
     let vec = vec.unwrap();
 
+    // PgVector is not cloneable, but in some cases we need a second copy of it
+    // for the insert.  This is a bit of a hack to get that second copy.
+    let spare_vec = LabeledVector::from_datums(values, isnull, &meta_page).unwrap();
+
     let heap_pointer = ItemPointer::with_item_pointer_data(*heap_tid);
     let mut storage = meta_page.get_storage_type();
     let mut stats = InsertStats::new();
@@ -189,10 +193,12 @@ unsafe fn aminsert_internal(
                 &heap_relation,
                 meta_page.get_distance_function(),
             );
+            assert!(vec.labels().is_none());
             insert_storage(
                 &plain,
                 &index_relation,
                 vec,
+                spare_vec,
                 heap_pointer,
                 &mut meta_page,
                 &mut stats,
@@ -209,6 +215,7 @@ unsafe fn aminsert_internal(
                 &bq,
                 &index_relation,
                 vec,
+                spare_vec,
                 heap_pointer,
                 &mut meta_page,
                 &mut stats,
@@ -222,6 +229,7 @@ unsafe fn insert_storage<S: Storage>(
     storage: &S,
     index_relation: &PgRelation,
     vector: LabeledVector,
+    spare_vector: LabeledVector,
     heap_pointer: ItemPointer,
     meta_page: &mut MetaPage,
     stats: &mut InsertStats,
@@ -237,7 +245,14 @@ unsafe fn insert_storage<S: Storage>(
     );
 
     let mut graph = Graph::new(GraphNeighborStore::Disk, meta_page);
-    graph.insert(index_relation, index_pointer, vector, storage, stats);
+    graph.insert(
+        index_relation,
+        index_pointer,
+        vector,
+        spare_vector,
+        storage,
+        stats,
+    );
 }
 
 #[pg_guard]
@@ -350,14 +365,14 @@ fn finalize_index_build<S: Storage>(
 ) -> usize {
     match state.graph.get_neighbor_store() {
         GraphNeighborStore::Builder(builder) => {
-            for (&index_pointer, neighbors) in builder.iter() {
+            for (&index_pointer, (labels, neighbors)) in builder.iter() {
                 write_stats.num_nodes += 1;
                 let prune_neighbors;
                 let neighbors =
                     if neighbors.len() > state.graph.get_meta_page().get_num_neighbors() as _ {
                         //OPT: get rid of this clone
                         prune_neighbors = state.graph.prune_neighbors(
-                            index_pointer,
+                            labels,
                             neighbors.clone(),
                             storage,
                             &mut write_stats.prune_stats,
@@ -459,13 +474,33 @@ unsafe extern "C" fn build_callback(
         StorageBuildState::SbqSpeedup(bq, state) => {
             let vec = LabeledVector::from_datums(values, isnull, state.graph.get_meta_page());
             if let Some(vec) = vec {
-                build_callback_memory_wrapper(&index_relation, heap_pointer, vec, state, *bq);
+                let spare_vec =
+                    LabeledVector::from_datums(values, isnull, state.graph.get_meta_page())
+                        .unwrap();
+                build_callback_memory_wrapper(
+                    &index_relation,
+                    heap_pointer,
+                    vec,
+                    spare_vec,
+                    state,
+                    *bq,
+                );
             }
         }
         StorageBuildState::Plain(plain, state) => {
             let vec = LabeledVector::from_datums(values, isnull, state.graph.get_meta_page());
             if let Some(vec) = vec {
-                build_callback_memory_wrapper(&index_relation, heap_pointer, vec, state, *plain);
+                let spare_vec =
+                    LabeledVector::from_datums(values, isnull, state.graph.get_meta_page())
+                        .unwrap();
+                build_callback_memory_wrapper(
+                    &index_relation,
+                    heap_pointer,
+                    vec,
+                    spare_vec,
+                    state,
+                    *plain,
+                );
             }
         }
     }
@@ -476,12 +511,13 @@ unsafe fn build_callback_memory_wrapper<S: Storage>(
     index: &PgRelation,
     heap_pointer: ItemPointer,
     vector: LabeledVector,
+    spare_vector: LabeledVector,
     state: &mut BuildState,
     storage: &mut S,
 ) {
     let mut old_context = state.memcxt.set_as_current();
 
-    build_callback_internal(index, heap_pointer, vector, state, storage);
+    build_callback_internal(index, heap_pointer, vector, spare_vector, state, storage);
 
     old_context.set_as_current();
     state.memcxt.reset();
@@ -492,6 +528,7 @@ fn build_callback_internal<S: Storage>(
     index: &PgRelation,
     heap_pointer: ItemPointer,
     vector: LabeledVector,
+    spare_vector: LabeledVector,
     state: &mut BuildState,
     storage: &mut S,
 ) {
@@ -520,9 +557,14 @@ fn build_callback_internal<S: Storage>(
         &mut state.stats,
     );
 
-    state
-        .graph
-        .insert(index, index_pointer, vector, storage, &mut state.stats);
+    state.graph.insert(
+        index,
+        index_pointer,
+        vector,
+        spare_vector,
+        storage,
+        &mut state.stats,
+    );
 }
 
 const BUILD_PHASE_TRAINING: i64 = 0;

pgvectorscale/src/access_method/graph.rs

@@ -202,6 +202,7 @@ impl<'a> Graph<'a> {
         &mut self,
         storage: &S,
         neighbors_of: ItemPointer,
+        labels: Option<&LabelSet>,
         additional_neighbors: Vec<NeighborWithDistance>,
         stats: &mut PruneNeighborStats,
     ) -> (bool, Vec<NeighborWithDistance>) {
@@ -237,7 +238,7 @@ impl<'a> Graph<'a> {
 
         let (pruned, new_neighbors) =
             if candidates.len() > self.neighbor_store.max_neighbors(self.get_meta_page()) {
-                let new_list = self.prune_neighbors(neighbors_of, candidates, storage, stats);
+                let new_list = self.prune_neighbors(labels, candidates, storage, stats);
                 (true, new_list)
             } else {
                 (false, candidates)
@@ -247,6 +248,7 @@ impl<'a> Graph<'a> {
             storage,
             self.meta_page,
             neighbors_of,
+            labels.cloned(),
             new_neighbors.clone(),
             stats,
         );
@@ -278,6 +280,7 @@ impl<'a> Graph<'a> {
         &self,
         index_pointer: IndexPointer,
         query: LabeledVector,
+        no_filter: bool,
         storage: &S,
         stats: &mut GreedySearchStats,
     ) -> HashSet<NeighborWithDistance> {
@@ -286,7 +289,13 @@ impl<'a> Graph<'a> {
             //no nodes in the graph
             return HashSet::with_capacity(0);
         }
-        let start_nodes = start_nodes.unwrap().get_for_node(query.labels());
+
+        let start_nodes = if no_filter {
+            start_nodes.unwrap().get_for_node(None)
+        } else {
+            start_nodes.unwrap().get_for_node(query.labels())
+        };
+
         let dm = storage.get_query_distance_measure(query);
         let search_list_size = self.meta_page.get_search_list_size_for_build() as usize;
         let mut l = ListSearchResult::new(
@@ -299,7 +308,13 @@ impl<'a> Graph<'a> {
             storage,
         );
         let mut visited_nodes = HashSet::with_capacity(search_list_size);
-        self.greedy_search_iterate(&mut l, search_list_size, Some(&mut visited_nodes), storage);
+        self.greedy_search_iterate(
+            &mut l,
+            search_list_size,
+            no_filter,
+            Some(&mut visited_nodes),
+            storage,
+        );
         stats.combine(&l.stats);
         visited_nodes
     }
@@ -336,6 +351,7 @@ impl<'a> Graph<'a> {
         &self,
         lsr: &mut ListSearchResult<S::QueryDistanceMeasure, S::LSNPrivateData>,
         visit_n_closest: usize,
+        no_filter: bool,
         mut visited_nodes: Option<&mut HashSet<NeighborWithDistance>>,
         storage: &S,
     ) {
@@ -352,7 +368,7 @@ impl<'a> Graph<'a> {
                 }
             }
             lsr.stats.record_visit();
-            storage.visit_lsn(lsr, list_search_entry_idx, &self.neighbor_store);
+            storage.visit_lsn(lsr, list_search_entry_idx, &self.neighbor_store, no_filter);
         }
     }
 
@@ -363,7 +379,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,
+        labels: Option<&LabelSet>,
         mut candidates: Vec<NeighborWithDistance>,
         storage: &S,
         stats: &mut PruneNeighborStats,
@@ -393,6 +409,7 @@ impl<'a> Graph<'a> {
                 if results.len() >= self.get_meta_page().get_num_neighbors() as _ {
                     return results;
                 }
+
                 if max_factors[i] > alpha {
                     continue;
                 }
@@ -412,13 +429,27 @@ impl<'a> Graph<'a> {
                     )
                 };
 
+                // TODO: optimization: precompute intersection of `labels` and `existing_neighbor.get_labels()`
+                // and use this for the `contains_intersection` check inside the loop
+
                 //go thru the other candidates (tail of the list)
                 for (j, candidate_neighbor) in candidates.iter().enumerate().skip(i + 1) {
                     //has it been completely excluded?
                     if max_factors[j] > max_alpha {
                         continue;
                     }
 
+                    // Does it contain essential labels?
+                    if let Some(labels) = labels {
+                        if !existing_neighbor
+                            .get_labels()
+                            .unwrap()
+                            .contains_intersection(candidate_neighbor.get_labels().unwrap(), labels)
+                        {
+                            continue;
+                        }
+                    }
+
                     let raw_distance_between_candidate_and_existing_neighbor = unsafe {
                         dist_state
                             .get_distance(candidate_neighbor.get_index_pointer_to_neighbor(), stats)
@@ -468,6 +499,7 @@ impl<'a> Graph<'a> {
                     storage,
                     self.meta_page,
                     index_pointer,
+                    vec.labels().cloned(),
                     Vec::<NeighborWithDistance>::with_capacity(
                         self.neighbor_store.max_neighbors(self.meta_page) as _,
                     ),
@@ -494,12 +526,11 @@ impl<'a> Graph<'a> {
 
     /// Check that all nodes of the graph are reachable from the start node(s)
     #[allow(dead_code)]
-    pub fn debug_check_consistency<S: Storage>(
-        &mut self,
+    pub fn debug_count_reachable_nodes<S: Storage>(
+        &self,
         storage: &S,
         stats: &mut InsertStats,
-    ) -> bool {
-        let num_nodes = 0; // TODO: calculate from base table?
+    ) -> usize {
         if let Some(start_nodes) = self.meta_page.get_start_nodes() {
             let mut visited = HashSet::new();
             let mut to_visit = start_nodes.get_all_nodes();
@@ -521,9 +552,9 @@ impl<'a> Graph<'a> {
                     to_visit.push(ip);
                 }
             }
-            return visited.len() == num_nodes;
+            return visited.len();
         }
-        num_nodes == 0
+        0
     }
 
     fn debug_format_labels(labels: Option<LabelSet>) -> String {
@@ -612,24 +643,44 @@ digraph G {
         index: &PgRelation,
         index_pointer: IndexPointer,
         vec: LabeledVector,
+        spare_vec: LabeledVector,
         storage: &S,
         stats: &mut InsertStats,
     ) {
         self.update_start_nodes(index, index_pointer, &vec, storage, stats);
 
+        if vec.labels().is_some() {
+            // Insert starting from label start nodes and apply label filtering
+            self.insert_internal(index_pointer, spare_vec, false, storage, stats);
+        }
+
+        // Insert starting from default start node and avoid label filtering
+        self.insert_internal(index_pointer, vec, true, storage, stats);
+    }
+
+    fn insert_internal<S: Storage>(
+        &mut self,
+        index_pointer: IndexPointer,
+        vec: LabeledVector,
+        no_filter: bool,
+        storage: &S,
+        stats: &mut InsertStats,
+    ) {
         let labels = vec.labels().cloned();
 
         #[allow(clippy::mutable_key_type)]
         let v = self.greedy_search_for_build(
             index_pointer,
             vec,
+            no_filter,
             storage,
             &mut stats.greedy_search_stats,
         );
 
         let (_, neighbor_list) = self.add_neighbors(
             storage,
             index_pointer,
+            labels.as_ref(),
             v.into_iter().collect(),
             &mut stats.prune_neighbor_stats,
         );
@@ -641,6 +692,7 @@ digraph G {
             let neighbor_contains_new_point = self.update_back_pointer(
                 neighbor.get_index_pointer_to_neighbor(),
                 index_pointer,
+                neighbor.get_labels(),
                 labels.as_ref(),
                 neighbor.get_distance_with_tie_break(),
                 storage,
@@ -664,10 +716,12 @@ digraph G {
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
     fn update_back_pointer<S: Storage>(
         &mut self,
         from: IndexPointer,
         to: IndexPointer,
+        from_labels: Option<&LabelSet>,
         to_labels: Option<&LabelSet>,
         distance_with_tie_break: &DistanceWithTieBreak,
         storage: &S,
@@ -678,7 +732,7 @@ digraph G {
             distance_with_tie_break.clone(),
             to_labels.cloned(),
         )];
-        let (_pruned, n) = self.add_neighbors(storage, from, new.clone(), prune_stats);
+        let (_pruned, n) = self.add_neighbors(storage, from, from_labels, new.clone(), prune_stats);
         n.contains(&new[0])
     }
 }