result: store vector samples when requested

vmg · vmg · commit db63ddde4b25 · 2024-03-13T13:31:02.000+01:00
Signed-off-by: Vicent Marti &lt;vmg@strn.cat&gt;
diff --git a/AnnService/inc/Core/Common/QueryResultSet.h b/AnnService/inc/Core/Common/QueryResultSet.h
@@ -74,18 +74,29 @@ class QueryResultSet : public QueryResult
         return m_results[0].Dist;
     }
 
-    bool AddPoint(const SizeType index, float dist)
+    bool AddPoint(const SizeType index, float dist, const void *sample, SizeType sample_length)
     {
         if (dist < m_results[0].Dist || (dist == m_results[0].Dist && index < m_results[0].VID))
         {
             m_results[0].VID = index;
             m_results[0].Dist = dist;
+
+            if (m_withSamples && sample) {
+                m_results[0].Sample = ByteArray::Alloc(sample_length);
+                memcpy(m_results[0].Sample.Data(), sample, sample_length);
+            }
+
             Heapify(m_resultNum);
             return true;
         }
         return false;
     }
 
+    bool AddPoint(const SizeType index, float dist)
+    {
+        return AddPoint(index, dist, nullptr, 0);
+    }
+
     inline void SortResult()
     {
         for (int i = m_resultNum - 1; i >= 0; i--)
diff --git a/AnnService/inc/Core/SPANN/ExtraDynamicSearcher.h b/AnnService/inc/Core/SPANN/ExtraDynamicSearcher.h
@@ -1245,14 +1245,17 @@ namespace SPTAG::SPANN {
             SearchStats* p_stats, std::set<int>* truth, std::map<int, std::set<int>>* found) override
         {
             auto exStart = std::chrono::high_resolution_clock::now();
+            std::chrono::microseconds remainLimit = std::chrono::microseconds::max();
 
             // const auto postingListCount = static_cast<uint32_t>(p_exWorkSpace->m_postingIDs.size());
 
             p_exWorkSpace->m_deduper.clear();
 
-            auto exSetUpEnd = std::chrono::high_resolution_clock::now();
-
-            p_stats->m_exSetUpLatency = ((double)std::chrono::duration_cast<std::chrono::microseconds>(exSetUpEnd - exStart).count()) / 1000;
+            if (p_stats) {
+                auto exSetUpEnd = std::chrono::high_resolution_clock::now();
+                p_stats->m_exSetUpLatency = ((double)std::chrono::duration_cast<std::chrono::microseconds>(exSetUpEnd - exStart).count()) / 1000;
+                remainLimit = m_hardLatencyLimit - std::chrono::microseconds((int)p_stats->m_totalLatency);
+            }
 
             COMMON::QueryResultSet<ValueType>& queryResults = *((COMMON::QueryResultSet<ValueType>*) & p_queryResults);
 
@@ -1265,8 +1268,6 @@ namespace SPTAG::SPANN {
 
             std::vector<std::string> postingLists;
 
-            std::chrono::microseconds remainLimit = m_hardLatencyLimit - std::chrono::microseconds((int)p_stats->m_totalLatency);
-
             auto readStart = std::chrono::high_resolution_clock::now();
             db->MultiGet(p_exWorkSpace->m_postingIDs, &postingLists, remainLimit);
             auto readEnd = std::chrono::high_resolution_clock::now();
@@ -1302,7 +1303,7 @@ namespace SPTAG::SPANN {
                         continue;
                     }
                     auto distance2leaf = p_index->ComputeDistance(queryResults.GetQuantizedTarget(), vectorInfo + m_metaDataSize);
-                    queryResults.AddPoint(vectorID, distance2leaf);
+                    queryResults.AddPoint(vectorID, distance2leaf, vectorInfo + m_metaDataSize, m_vectorInfoSize - m_metaDataSize);
                 }
                 auto compEnd = std::chrono::high_resolution_clock::now();
                 if (realNum <= m_mergeThreshold && !m_opt->m_inPlace) MergeAsync(p_index.get(), curPostingID);
diff --git a/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h b/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h
@@ -126,7 +126,7 @@ namespace SPTAG
             if (p_exWorkSpace->m_deduper.CheckAndSet(vectorID)) continue; \
             (this->*m_parseEncoding)(p_index, listInfo, (ValueType*)(p_postingListFullData + offsetVector));\
             auto distance2leaf = p_index->ComputeDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); \
-            queryResults.AddPoint(vectorID, distance2leaf); \
+            queryResults.AddPoint(vectorID, distance2leaf, p_postingListFullData + offsetVector, p_index->GetFeatureDim() * sizeof(ValueType)); \
         } \
 
 #define ProcessPostingOffset() \
@@ -138,7 +138,7 @@ namespace SPTAG
             if (p_exWorkSpace->m_deduper.CheckAndSet(vectorID)) continue; \
             (this->*m_parseEncoding)(p_index, listInfo, (ValueType*)(p_postingListFullData + offsetVector));\
             auto distance2leaf = p_index->ComputeDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); \
-            queryResults.AddPoint(vectorID, distance2leaf); \
+            queryResults.AddPoint(vectorID, distance2leaf, p_postingListFullData + offsetVector, p_index->GetFeatureDim() * sizeof(ValueType)); \
             foundResult = true;\
             break;\
         } \
diff --git a/AnnService/inc/Core/SearchQuery.h b/AnnService/inc/Core/SearchQuery.h
@@ -23,7 +23,8 @@ class QueryResult
           m_quantizedTarget(nullptr),
           m_quantizedSize(0),
           m_resultNum(0),
-          m_withMeta(false)
+          m_withMeta(false),
+          m_withSamples(false)
     {
     }
 
@@ -32,9 +33,9 @@ class QueryResult
         Init(nullptr, p_resultNum, true);
     }
 
-    QueryResult(const void* p_target, int p_resultNum, bool p_withMeta)
+    QueryResult(const void* p_target, int p_resultNum, bool p_withMeta = false, bool p_withSamples = false)
     {
-        Init(p_target, p_resultNum, p_withMeta);
+        Init(p_target, p_resultNum, p_withMeta, p_withSamples);
     }
 
     
@@ -93,11 +94,12 @@ class QueryResult
     }
 
 
-    inline void Init(const void* p_target, int p_resultNum, bool p_withMeta)
+    inline void Init(const void* p_target, int p_resultNum, bool p_withMeta = false, bool p_withSamples = false)
     {
         m_target = p_target;
         m_resultNum = p_resultNum;
         m_withMeta = p_withMeta;
+        m_withSamples = p_withSamples;
         m_quantizedTarget = (void*)p_target;
         m_quantizedSize = 0;
 
@@ -177,6 +179,11 @@ class QueryResult
         return m_results.Data();
     }
 
+    inline bool WithSamples() const
+    {
+        return m_withSamples;
+    }
+
 
     inline bool WithMeta() const
     {
@@ -211,6 +218,7 @@ class QueryResult
             m_results[i].VID = -1;
             m_results[i].Dist = MaxDist;
             m_results[i].Meta.Clear();
+            m_results[i].Sample.Clear();
         }
     }
 
@@ -249,6 +257,7 @@ class QueryResult
     int m_resultNum;
 
     bool m_withMeta;
+    bool m_withSamples;
 
     Array<BasicResult> m_results;
 };
diff --git a/AnnService/inc/Core/SearchResult.h b/AnnService/inc/Core/SearchResult.h
@@ -67,6 +67,7 @@ namespace SPTAG
         SizeType VID;
         float Dist;
         ByteArray Meta;
+        ByteArray Sample;
         bool RelaxedMono;
 
         BasicResult() : VID(-1), Dist(MaxDist), RelaxedMono(false) {}
diff --git a/AnnService/src/Core/BKT/BKTIndex.cpp b/AnnService/src/Core/BKT/BKTIndex.cpp
@@ -316,7 +316,7 @@ namespace SPTAG
                         {
                             if (checkFilter(m_pMetadata, tmpNode, filterFunc))
                             {
-                                p_query.AddPoint(tmpNode, gnode.distance);
+                                p_query.AddPoint(tmpNode, gnode.distance, (m_pSamples)[tmpNode], GetFeatureDim() * sizeof(T));
                             }
                         }
                     }
@@ -379,7 +379,7 @@ namespace SPTAG
                 }
                 if (notDeleted(m_deletedID, tmpNode))
                 {
-                    p_query.AddPoint(tmpNode, gnode.distance);
+                    p_query.AddPoint(tmpNode, gnode.distance, (m_pSamples)[tmpNode], GetFeatureDim() * sizeof(T));
                     count++;
                     if (gnode.distance > p_space.m_Results.worst() || 
                         p_space.m_iNumberOfCheckedLeaves > p_space.m_iMaxCheck) {
diff --git a/AnnService/src/Core/KDT/KDTIndex.cpp b/AnnService/src/Core/KDT/KDTIndex.cpp
@@ -200,7 +200,7 @@ namespace SPTAG
                     
                 if (notDeleted(m_deletedID, gnode.node))
                 {
-                    if (!p_query.AddPoint(gnode.node, gnode.distance) && p_space.m_iNumberOfCheckedLeaves > p_space.m_iMaxCheck) 
+                    if (!p_query.AddPoint(gnode.node, gnode.distance, (m_pSamples)[gnode.node], GetFeatureDim() * sizeof(T)) && p_space.m_iNumberOfCheckedLeaves > p_space.m_iMaxCheck) 
                     {
                         p_query.SortResult();
                         return;

Original file line number	Diff line number	Diff line change
`@@ -74,18 +74,29 @@ class QueryResultSet : public QueryResult`
`74`	`74`	`return m_results[0].Dist;`
`75`	`75`	`}`
`76`	`76`
`77`		`- bool AddPoint(const SizeType index, float dist)`
	`77`	`+ bool AddPoint(const SizeType index, float dist, const void *sample, SizeType sample_length)`
`78`	`78`	`{`
`79`	`79`	`if (dist < m_results[0].Dist \|\| (dist == m_results[0].Dist && index < m_results[0].VID))`
`80`	`80`	`{`
`81`	`81`	`m_results[0].VID = index;`
`82`	`82`	`m_results[0].Dist = dist;`
	`83`	`+`
	`84`	`+ if (m_withSamples && sample) {`
	`85`	`+ m_results[0].Sample = ByteArray::Alloc(sample_length);`
	`86`	`+ memcpy(m_results[0].Sample.Data(), sample, sample_length);`
	`87`	`+ }`
	`88`	`+`
`83`	`89`	`Heapify(m_resultNum);`
`84`	`90`	`return true;`
`85`	`91`	`}`
`86`	`92`	`return false;`
`87`	`93`	`}`
`88`	`94`
	`95`	`+ bool AddPoint(const SizeType index, float dist)`
	`96`	`+ {`
	`97`	`+ return AddPoint(index, dist, nullptr, 0);`
	`98`	`+ }`
	`99`	`+`
`89`	`100`	`inline void SortResult()`
`90`	`101`	`{`
`91`	`102`	`for (int i = m_resultNum - 1; i >= 0; i--)`
Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,8 @@ class QueryResult`
`23`	`23`	`m_quantizedTarget(nullptr),`
`24`	`24`	`m_quantizedSize(0),`
`25`	`25`	`m_resultNum(0),`
`26`		`- m_withMeta(false)`
	`26`	`+ m_withMeta(false),`
	`27`	`+ m_withSamples(false)`
`27`	`28`	`{`
`28`	`29`	`}`
`29`	`30`
`@@ -32,9 +33,9 @@ class QueryResult`
`32`	`33`	`Init(nullptr, p_resultNum, true);`
`33`	`34`	`}`
`34`	`35`
`35`		`- QueryResult(const void* p_target, int p_resultNum, bool p_withMeta)`
	`36`	`+ QueryResult(const void* p_target, int p_resultNum, bool p_withMeta = false, bool p_withSamples = false)`
`36`	`37`	`{`
`37`		`- Init(p_target, p_resultNum, p_withMeta);`
	`38`	`+ Init(p_target, p_resultNum, p_withMeta, p_withSamples);`
`38`	`39`	`}`
`39`	`40`
`40`	`41`
`@@ -93,11 +94,12 @@ class QueryResult`
`93`	`94`	`}`
`94`	`95`
`95`	`96`
`96`		`- inline void Init(const void* p_target, int p_resultNum, bool p_withMeta)`
	`97`	`+ inline void Init(const void* p_target, int p_resultNum, bool p_withMeta = false, bool p_withSamples = false)`
`97`	`98`	`{`
`98`	`99`	`m_target = p_target;`
`99`	`100`	`m_resultNum = p_resultNum;`
`100`	`101`	`m_withMeta = p_withMeta;`
	`102`	`+ m_withSamples = p_withSamples;`
`101`	`103`	`m_quantizedTarget = (void*)p_target;`
`102`	`104`	`m_quantizedSize = 0;`
`103`	`105`
`@@ -177,6 +179,11 @@ class QueryResult`
`177`	`179`	`return m_results.Data();`
`178`	`180`	`}`
`179`	`181`
	`182`	`+ inline bool WithSamples() const`
	`183`	`+ {`
	`184`	`+ return m_withSamples;`
	`185`	`+ }`
	`186`	`+`
`180`	`187`
`181`	`188`	`inline bool WithMeta() const`
`182`	`189`	`{`
`@@ -211,6 +218,7 @@ class QueryResult`
`211`	`218`	`m_results[i].VID = -1;`
`212`	`219`	`m_results[i].Dist = MaxDist;`
`213`	`220`	`m_results[i].Meta.Clear();`
	`221`	`+ m_results[i].Sample.Clear();`
`214`	`222`	`}`
`215`	`223`	`}`
`216`	`224`
`@@ -249,6 +257,7 @@ class QueryResult`
`249`	`257`	`int m_resultNum;`
`250`	`258`
`251`	`259`	`bool m_withMeta;`
	`260`	`+ bool m_withSamples;`
`252`	`261`
`253`	`262`	`Array<BasicResult> m_results;`
`254`	`263`	`};`
Original file line number	Diff line number	Diff line change
`@@ -316,7 +316,7 @@ namespace SPTAG`
`316`	`316`	`{`
`317`	`317`	`if (checkFilter(m_pMetadata, tmpNode, filterFunc))`
`318`	`318`	`{`
`319`		`- p_query.AddPoint(tmpNode, gnode.distance);`
	`319`	`+ p_query.AddPoint(tmpNode, gnode.distance, (m_pSamples)[tmpNode], GetFeatureDim() * sizeof(T));`
`320`	`320`	`}`
`321`	`321`	`}`
`322`	`322`	`}`
`@@ -379,7 +379,7 @@ namespace SPTAG`
`379`	`379`	`}`
`380`	`380`	`if (notDeleted(m_deletedID, tmpNode))`
`381`	`381`	`{`
`382`		`- p_query.AddPoint(tmpNode, gnode.distance);`
	`382`	`+ p_query.AddPoint(tmpNode, gnode.distance, (m_pSamples)[tmpNode], GetFeatureDim() * sizeof(T));`
`383`	`383`	`count++;`
`384`	`384`	`if (gnode.distance > p_space.m_Results.worst() \|\|`
`385`	`385`	`p_space.m_iNumberOfCheckedLeaves > p_space.m_iMaxCheck) {`
Original file line number	Diff line number	Diff line change
`@@ -200,7 +200,7 @@ namespace SPTAG`
`200`	`200`
`201`	`201`	`if (notDeleted(m_deletedID, gnode.node))`
`202`	`202`	`{`
`203`		`- if (!p_query.AddPoint(gnode.node, gnode.distance) && p_space.m_iNumberOfCheckedLeaves > p_space.m_iMaxCheck)`
	`203`	`+ if (!p_query.AddPoint(gnode.node, gnode.distance, (m_pSamples)[gnode.node], GetFeatureDim() * sizeof(T)) && p_space.m_iNumberOfCheckedLeaves > p_space.m_iMaxCheck)`
`204`	`204`	`{`
`205`	`205`	`p_query.SortResult();`
`206`	`206`	`return;`