feat(cpp/search): return std::vector<ResultSet> from rerank and optimize heap updates

- Refactor deglib::search::rerank signature to return std::vector<ResultSet>
- Add PQV::replace_top returning const ObjectType& for O(log k) single-pass sift-down root replacement
- Avoid heap.top() overhead in rerank loop by tracking max_dist locally
- Add docstrings documenting unsorted heap ordering contract of ResultSet
- Update Python bindings (float_space_rerank) and FloatSpace.rerank API with return_distances support
- Update C++ unit tests in test_search.cpp

Author: Neiko2002

cpp/deglib/include/deglib/graph/internal_graph.h

@@ -87,20 +87,53 @@ class PQV : public std::vector<ObjectType> {
         this->pop_back();
     }
 
-    /// Replaces the top element in O(log k) time without vector reallocation.
-    void replace_top(ObjectType&& x) {
-        std::pop_heap(this->begin(), this->end(), comp);
-        this->back() = std::move(x);
-        std::push_heap(this->begin(), this->end(), comp);
+    /// Replaces the top element in-place, sifts it down in single-pass O(log k) time,
+    /// and returns a const reference to the new top element.
+    template <class... Args>
+    const ObjectType& replace_top(Args&&... args) {
+        ObjectType val(std::forward<Args>(args)...);
+        size_t len = this->size();
+        size_t parent = 0;
+        size_t child = 1;
+
+        while (child < len) {
+            // Find the larger child (comp(a, b) means a < b for max-heap)
+            if (child + 1 < len && comp((*this)[child], (*this)[child + 1])) {
+                child++;
+            }
+            // If val is not less than the larger child, heap property holds
+            if (!comp(val, (*this)[child])) {
+                break;
+            }
+            (*this)[parent] = std::move((*this)[child]);
+            parent = child;
+            child = 2 * parent + 1;
+        }
+        (*this)[parent] = std::move(val);
+        return this->front();
     }
 
     /// Re-establishes the heap order using the internal comparator.
     void heapify() {
         std::make_heap(this->begin(), this->end(), comp);
     }
+
+    /// Sorts the heap elements in ascending order using the internal comparator,
+    /// consuming the heap property. After calling this, top() is no longer valid
+    /// until heapify() is called again.
+    void sort() {
+        std::sort_heap(this->begin(), this->end(), comp);
+    }
 };
 
-// Search result set (Max-Heap: top() returns the entry with the LARGEST distance)
+/**
+ * Search result set (Max-Heap: top() returns the entry with the LARGEST distance).
+ * 
+ * Note: Elements are maintained in max-heap order within the underlying vector.
+ * Iterating or indexing the container directly will yield elements in arbitrary (unsorted) heap order.
+ * To access elements in ordered sequence (closest distance first), extract them using top() and pop(),
+ * or call sort() explicitly to sort the underlying storage in ascending distance order.
+ */
 using ResultSet = PQV<std::less<ObjectDistance>, ObjectDistance>;
 
 // Unchecked search candidates (Min-Heap: top() returns the entry with the SMALLEST distance)

cpp/deglib/include/deglib/graph/visited_list_pool.h

@@ -5,6 +5,7 @@
 #include <vector>
 #include <deque>
 #include <cstdint>
+#include <memory>
 
 /**
  * Ref https://raw.githubusercontent.com/nmslib/hnswlib/master/hnswlib/visited_list_pool.h

cpp/deglib/include/deglib/search.h

@@ -1,10 +1,13 @@
 #pragma once
 
 #include <cstddef>
+#include <cstdint>
+#include <limits>
 #include <queue>
 #include <span>
 #include <stdexcept>
 #include <string>
+#include <vector>
 #include "deglib/distances.h"
 #include "deglib/filter.h"
 #include "deglib/concurrent.h"
@@ -18,31 +21,34 @@ using ResultSet = deglib::graph::ResultSet;
 /**
  * Rerank candidate neighbor indices for queries using exact FloatSpace distances.
  *
+ * Returns a vector of ResultSet objects (one per query). Each ResultSet contains the top-k nearest
+ * candidates maintained in a max-heap property. Note that ResultSet elements are in heap order (unsorted);
+ * use top() and pop(), or call sort() on the ResultSet if sorted order is required.
+ *
  * @param space                FloatSpace distance calculator instance
  * @param queries              Pointer to [num_queries x dim] query vectors
  * @param num_queries          Number of query vectors
  * @param base_vectors         Pointer to [num_base_vectors x dim] target/base vectors (if null, queries are used as targets)
  * @param num_base_vectors     Number of base vectors
  * @param base_candidates      Pointer to [num_queries x candidates_per_query] candidate indices
  * @param candidates_per_query Number of candidate indices provided per query
- * @param k_top                Number of top candidates to output per query (0 = all)
+ * @param k_top                Number of top candidates to keep per query (0 = all)
  * @param num_threads          Number of worker threads (0 = auto-detect)
- * @param out_result_indices   Output pointer to [num_queries x k_top] uint32_t indices
+ * @return std::vector<ResultSet> containing the top-k result sets per query (unsorted heap order)
  */
-inline void rerank(
+inline std::vector<ResultSet> rerank(
     const deglib::distances::FloatSpace& space,
     const void* queries,
     size_t num_queries,
     const void* base_vectors,
     size_t num_base_vectors,
     const uint32_t* base_candidates,
     size_t candidates_per_query,
-    size_t k_top,
-    size_t num_threads,
-    uint32_t* out_result_indices
+    size_t k_top = 0,
+    size_t num_threads = 0
 ) {
-    if (queries == nullptr || base_candidates == nullptr || out_result_indices == nullptr) {
-        throw std::invalid_argument("rerank: queries, base_candidates, and out_result_indices must not be null");
+    if (queries == nullptr || base_candidates == nullptr) {
+        throw std::invalid_argument("rerank: queries and base_candidates must not be null");
     }
 
     if (k_top == 0 || k_top > candidates_per_query) {
@@ -59,15 +65,18 @@ inline void rerank(
     const uint8_t* t_ptr = static_cast<const uint8_t*>(target_vectors);
 
     const auto param = space.get_dist_func_param();
-    // Resolve variant type ONCE outside loops via compile-time static dispatch
+
+    std::vector<ResultSet> results(num_queries);
+
     space.compute([&](const auto& dist_func_obj) {
         using DistType = std::decay_t<decltype(dist_func_obj)>;
         deglib::concurrent::parallel_for(0, num_queries, num_threads, [&](size_t i, size_t) {
             const uint8_t* query_ptr = q_ptr + i * byte_stride_query;
             const uint32_t* cand_row = base_candidates + i * candidates_per_query;
 
-            std::vector<std::pair<float, uint32_t>> heap;
-            heap.reserve(k_top + 1);
+            ResultSet heap;
+            heap.reserve(k_top);
+            float max_dist = std::numeric_limits<float>::max();
 
             for (size_t j = 0; j < candidates_per_query; ++j) {
                 uint32_t cand_idx = cand_row[j];
@@ -78,29 +87,26 @@ inline void rerank(
                 const uint8_t* cand_ptr = t_ptr + cand_idx * byte_stride_target;
                 float dist = DistType::compare(query_ptr, cand_ptr, param);
 
+                // Keep candidates in max-heap of capacity k_top.
+                // Until the heap reaches k_top elements, add all valid candidates.
                 if (heap.size() < k_top) {
-                    heap.push_back({dist, cand_idx});
-                    std::push_heap(heap.begin(), heap.end());
-                } else if (dist < heap.front().first) {
-                    std::pop_heap(heap.begin(), heap.end());
-                    heap.back() = {dist, cand_idx};
-                    std::push_heap(heap.begin(), heap.end());
+                    heap.emplace(cand_idx, dist);
+                    if (heap.size() == k_top) {
+                        max_dist = heap.top().getDistance();
+                    }
+                } 
+                // Once full, only consider candidates strictly closer than the worst (max_dist).
+                // replace_top replaces the root in O(log k) and returns a reference to the NEW root element (new max distance).
+                else if (dist < max_dist) {
+                    max_dist = heap.replace_top(cand_idx, dist).getDistance();
                 }
             }
 
-            std::sort_heap(heap.begin(), heap.end());
-
-            uint32_t* out_row = out_result_indices + i * k_top;
-            size_t actual_k = heap.size();
-            for (size_t k = 0; k < actual_k; ++k) {
-                out_row[k] = heap[k].second;
-            }
-            for (size_t k = actual_k; k < k_top; ++k) {
-                out_row[k] = static_cast<uint32_t>(i);
-            }
+            results[i] = std::move(heap);
         });
     });
+
+    return results;
 }
 
 } // namespace deglib::search
-

cpp/test/src/unit/graph/test_internal_graph.cpp

@@ -133,6 +133,71 @@ TEST(UncheckedSet, PopRemovesMin) {
    EXPECT_NEAR(us.top().getDistance(), 3.0f, 1e-6f);
 }
 
+// ---------------------------------------------------------------------------
+//  PQV emplace / pop / heapify (heap invariant tests)
+// ---------------------------------------------------------------------------
+
+TEST(PQV, EmplaceRestoresHeap) {
+   deglib::graph::ResultSet rs;
+   rs.emplace(1, 1.0f);
+   rs.emplace(2, 5.0f);
+   rs.emplace(3, 3.0f);
+
+   // ResultSet is a max-heap: top() must be the largest distance
+   EXPECT_EQ(rs.size(), 3u);
+   EXPECT_NEAR(rs.top().getDistance(), 5.0f, 1e-6f);
+   EXPECT_EQ(rs.top().getIdentifier(), 2u);
+
+   // Verify heap property holds for all elements
+   for (size_t i = 0; i < rs.size(); ++i) {
+      EXPECT_GE(rs[i].getDistance(), 0.0f);
+   }
+}
+
+TEST(PQV, Heapify) {
+   deglib::graph::ResultSet rs;
+   // Insert elements without maintaining heap order (bypass emplace)
+   rs.reserve(5);
+   static_cast<std::vector<deglib::graph::ObjectDistance>&>(rs).emplace_back(1, 1.0f);
+   static_cast<std::vector<deglib::graph::ObjectDistance>&>(rs).emplace_back(2, 5.0f);
+   static_cast<std::vector<deglib::graph::ObjectDistance>&>(rs).emplace_back(3, 3.0f);
+
+   // Before heapify, top() is front() which may not be the max
+   rs.heapify();
+   EXPECT_NEAR(rs.top().getDistance(), 5.0f, 1e-6f);
+   EXPECT_EQ(rs.top().getIdentifier(), 2u);
+}
+
+TEST(PQV, PopMaintainsHeap) {
+   deglib::graph::ResultSet rs;
+   rs.emplace(1, 1.0f);
+   rs.emplace(2, 5.0f);
+   rs.emplace(3, 3.0f);
+
+   rs.pop(); // remove max (5.0)
+   EXPECT_EQ(rs.size(), 2u);
+   EXPECT_NEAR(rs.top().getDistance(), 3.0f, 1e-6f);
+
+   rs.pop(); // remove max (3.0)
+   EXPECT_EQ(rs.size(), 1u);
+   EXPECT_NEAR(rs.top().getDistance(), 1.0f, 1e-6f);
+}
+
+// ---------------------------------------------------------------------------
+//  ObjectDistance POD traits
+// ---------------------------------------------------------------------------
+
+TEST(ObjectDistance, POD_Traits) {
+   // ObjectDistance should be trivially default-constructible, copyable, and movable
+   static_assert(std::is_trivially_default_constructible_v<deglib::graph::ObjectDistance>);
+   static_assert(std::is_trivially_copyable_v<deglib::graph::ObjectDistance>);
+   static_assert(std::is_trivially_destructible_v<deglib::graph::ObjectDistance>);
+   static_assert(std::is_standard_layout_v<deglib::graph::ObjectDistance>);
+
+   // Verify size: uint32_t + float = 8 bytes
+   EXPECT_EQ(sizeof(deglib::graph::ObjectDistance), 8u);
+}
+
 // ---------------------------------------------------------------------------
 //  Mock implementation of InternalGraph for testing the interface
 // ---------------------------------------------------------------------------