GetAllowLessThanKResults in context class

microsoft · Dec 22, 2023 · 8a3329e · 8a3329e
1 parent 77b432f
commit 8a3329e
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 4 deletions.
diff --git a/include/parameters.h b/include/parameters.h
@@ -139,8 +139,8 @@ enum State : uint8_t
 template <typename LabelT = uint32_t> class IndexSearchContext
 {
   public:
-    IndexSearchContext(uint32_t time_limit_in_microseconds = 0u, uint32_t io_limit = UINT32_MAX)
-        : _time_limit_in_microseconds(time_limit_in_microseconds), _io_limit(io_limit), _result_state(State::Unknown)
+    IndexSearchContext(uint32_t time_limit_in_microseconds = 0u, uint32_t io_limit = UINT32_MAX, bool allowLessThanKResults = false)
+        : _time_limit_in_microseconds(time_limit_in_microseconds), _io_limit(io_limit), _result_state(State::Unknown), _allowLessThankResults(allowLessThanKResults)
     {
         _use_filter = false;
         _label = (LabelT)0;
@@ -198,6 +198,11 @@ template <typename LabelT = uint32_t> class IndexSearchContext
         return _stats;
     }
 
+    bool GetAllowLessThanKResults()
+    {
+        return _allowLessThankResults;
+    }
+
   private:
     uint32_t _time_limit_in_microseconds;
     uint32_t _io_limit;
@@ -206,6 +211,7 @@ template <typename LabelT = uint32_t> class IndexSearchContext
     LabelT _label;
     Timer _timer;
     QueryStats _stats;
+    bool _allowLessThankResults;
 };
 
 } // namespace diskann
diff --git a/src/index.cpp b/src/index.cpp
@@ -2251,10 +2251,14 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::search(const T *query, con
             break;
     }
 
-    if (pos <= K)
+    if (pos <= K && context.GetAllowLessThanKResults())
     {
         context.SetState(State::Success);
     }
+    else if(pos < K)
+    {
+        context.SetState(State::Failure);
+    }
     else
     {
         context.SetState(State::Failure);
@@ -2386,10 +2390,14 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::search_with_filters(const
         if (pos == K)
             break;
     }
-    if (pos <= K)
+    if (pos <= K && context.GetAllowLessThanKResults())
     {
         context.SetState(State::Success);
     }
+    else if(pos < K)
+    {
+        context.SetState(State::Failure);
+    }
     else
     {
         context.SetState(State::Failure);