fix memory leak

microsoft · Mar 4, 2024 · d5207b7 · d5207b7
1 parent a25ee6f
commit d5207b7
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 0 deletions.
diff --git a/include/pq_scratch.h b/include/pq_scratch.h
@@ -17,6 +17,7 @@ template <typename T> class PQScratch
 
     PQScratch(size_t graph_degree, size_t aligned_dim);
     void initialize(size_t dim, const T *query, const float norm = 1.0f);
+    virtual ~PQScratch();
 };
 
 } // namespace diskann
diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp
@@ -93,6 +93,10 @@ template <typename T, typename LabelT> PQFlashIndex<T, LabelT>::~PQFlashIndex()
     {
         delete[] _pts_to_labels;
     }
+    if (_medoids != nullptr) 
+    {
+        delete[] _medoids;
+    }
 }
 
 template <typename T, typename LabelT> inline uint64_t PQFlashIndex<T, LabelT>::get_node_sector(uint64_t node_id)

diff --git a/src/scratch.cpp b/src/scratch.cpp
@@ -143,6 +143,15 @@ template <typename T> PQScratch<T>::PQScratch(size_t graph_degree, size_t aligne
     memset(rotated_query, 0, aligned_dim * sizeof(float));
 }
 
+template <typename T> PQScratch<T>::~PQScratch()
+{
+    diskann::aligned_free((void *)aligned_pq_coord_scratch);
+    diskann::aligned_free((void *)aligned_pqtable_dist_scratch);
+    diskann::aligned_free((void *)aligned_dist_scratch);
+    diskann::aligned_free((void *)aligned_query_float);
+    diskann::aligned_free((void *)rotated_query);
+}
+
 template <typename T> void PQScratch<T>::initialize(size_t dim, const T *query, const float norm)
 {
     for (size_t d = 0; d < dim; ++d)