Implement the HNSW drop and rollback logic

biigle · Jan 8, 2025 · 87ed509 · 87ed509
1 parent 11989cd
commit 87ed509
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 6 deletions.
diff --git a/app/Http/Controllers/Api/ImageAnnotationController.php b/app/Http/Controllers/Api/ImageAnnotationController.php
@@ -232,7 +232,7 @@ public function store(StoreImageAnnotation $request)
         if (is_null($labelId) && $request->has('feature_vector')) {
             // Add labelBOTlabels attribute to the response.
             $annotation->append('labelBOTLabels');
-            
+
             // Get label tree id(s).
             $trees = $this->getLabelTreeIds($request->user(), $image->volume_id);
 
@@ -437,6 +437,10 @@ protected function performVectorSearch($featureVector, $trees, $topNLabels)
     */
     protected function performAnnSearch($featureVector, $trees)
     {
+        // check if the HNSW index exists
+        if (!$this->indexExists(config('labelbot.HNSW_ImgAnno_index_name'))) {
+            return [];
+        }
         $subquery = ImageAnnotationLabelFeatureVector::select('label_id', 'label_tree_id')
             ->selectRaw('(vector <=> ?) AS distance', [$featureVector])
             ->orderBy('distance')
@@ -465,25 +469,54 @@ protected function performAnnSearch($featureVector, $trees)
     */
     protected function performKnnSearch($featureVector, $trees)
     {
+        // Drop HNSW index temporarily
+        DB::beginTransaction();
+        $this->dropHNSWIndex();
+
         $subquery = ImageAnnotationLabelFeatureVector::select('label_id', 'label_tree_id')
             ->selectRaw('(vector <=> ?) AS distance', [$featureVector])
             ->whereIn('label_tree_id', $trees) // Apply label tree ID filter in the subquery to use the B-Tree index for faster filtering
             ->orderBy('distance')
             ->limit(config('labelbot.K')); // K = 100
 
-        // TODO: Drop HNSW index temporary
-        // DB::beginTransaction();
-
+        // Save results.
         $topNLabels = DB::query()->fromSub($subquery, 'subquery')
             ->groupBy('label_id')
             ->orderByRaw('MIN(distance)')
             ->limit(config('labelbot.N')) // N = 3
             ->pluck('label_id')
             ->toArray();
 
-        // TODO: Rollback the HNSW index drop
-        // DB::rollback();
+        // Rollback the HNSW index drop
+        DB::rollback();
 
         return $topNLabels;
     }
+
+    /**
+     * Check if the index exists.
+     * 
+     * @param string $indexName The index name.
+     * 
+     * @return boolean
+     */
+    protected function indexExists($indexName)
+    {
+        return !empty(DB::select("SELECT indexname FROM pg_indexes WHERE indexname = '$indexName'"));
+    }
+
+    /**
+     * Drop the HNSW index if exists. This step is necessary to perform exact KNN search
+     * because the planner almost always prioritize the HNSW index to perform vector search.
+     */
+
+    protected function dropHNSWIndex()
+    {
+        // Check if the index exists
+        $indexName = config('labelbot.HNSW_ImgAnno_index_name');
+        if ($this->indexExists($indexName)) {
+            // Drop the index
+            DB::statement("DROP INDEX $indexName");
+        }
+    }
 }
diff --git a/config/labelbot.php b/config/labelbot.php
@@ -1,3 +1,4 @@
+
 <?php
 
 return [
@@ -24,4 +25,27 @@
     | and returned.
     */
     'N' => 3,
+
+    /*
+    |--------------------------------------------------------------------------
+    | HNSW Index name (Image Annotation)
+    |--------------------------------------------------------------------------
+    |
+    | This is the name used when creating, dropping or checking the HNSW index.
+    | The HNSW index is built on the vector column in the
+    | image_annotation_label_feature_vectors table.
+    */
+    'HNSW_ImgAnno_index_name' => 'image_annotation_label_feature_vectors_vector_idx',
+
+        /*
+    |--------------------------------------------------------------------------
+    | B-Tree Index name (Image Annotation)
+    |--------------------------------------------------------------------------
+    |
+    | This is the name used when creating, dropping or checking the B-Tree index.
+    | The B-Tree index is built on the label_tree_id column in the
+    | image_annotation_label_feature_vectors table.
+    */
+    'B_Tree_ImgAnno_index_name' => 'image_annotation_label_feature_vectors_label_tree_id_index',
+
 ];