Add match features to search result in vdslib module.

3 files changed, 67 insertions, 7 deletions

diff --git a/vdslib/src/tests/container/searchresulttest.cpp b/vdslib/src/tests/container/searchresulttest.cpp
index 8433a6008a1..2b27aeeb95c 100644
--- a/vdslib/src/tests/container/searchresulttest.cpp
+++ b/vdslib/src/tests/container/searchresulttest.cpp
@@ -1,11 +1,26 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
-#include <vespa/document/base/testdocman.h>
 #include <vespa/vdslib/container/searchresult.h>
 #include <vespa/vespalib/gtest/gtest.h>
 
+using vespalib::FeatureValues;
+using FeatureValue = vespalib::FeatureSet::Value;
+
 namespace vdslib {
 
+namespace {
+
+std::vector<double> to_doubles(vespalib::ConstArrayRef<FeatureValue> v) {
+    std::vector<double> result;
+    for (auto& iv : v) {
+        EXPECT_TRUE(iv.is_double());
+        result.emplace_back(iv.as_double());
+    }
+    return result;
+}
+
+}
+
 TEST(SearchResultTest, test_simple)
 {
     SearchResult a;
@@ -68,4 +83,31 @@ TEST(SearchResultTest, test_simple_sort_data)
     EXPECT_EQ(4, sz);
 }
 
+TEST(SearchResultTest, test_match_features)
+{
+    SearchResult sr;
+    sr.addHit(7, "doc1", 5);
+    sr.addHit(8, "doc2", 7);
+    FeatureValues mf;
+    mf.names.push_back("foo");
+    mf.names.push_back("bar");
+    mf.values.resize(4);
+    mf.values[0].set_double(1.0);
+    mf.values[1].set_double(7.0);
+    mf.values[2].set_double(12.0);
+    mf.values[3].set_double(13.0);
+    sr.set_match_features(FeatureValues(mf));
+    EXPECT_EQ(mf.names, sr.get_match_features().names);
+    EXPECT_EQ(mf.values, sr.get_match_features().values);
+    EXPECT_EQ((std::vector<double>{  1.0,  7.0}), to_doubles(sr.get_match_feature_values(0)));
+    EXPECT_EQ((std::vector<double>{ 12.0, 13.0}), to_doubles(sr.get_match_feature_values(1)));
+    sr.sort();
+    // Sorting does not change the stored match features
+    EXPECT_EQ(mf.names, sr.get_match_features().names);
+    EXPECT_EQ(mf.values, sr.get_match_features().values);
+    // Sorting affects retrieval of the stored matched features
+    EXPECT_EQ((std::vector<double>{ 12.0, 13.0}), to_doubles(sr.get_match_feature_values(0)));
+    EXPECT_EQ((std::vector<double>{  1.0,  7.0}), to_doubles(sr.get_match_feature_values(1)));
+}
+
 }
diff --git a/vdslib/src/vespa/vdslib/container/searchresult.cpp b/vdslib/src/vespa/vdslib/container/searchresult.cpp
index 3d7ae1631af..c110eaa792d 100644
--- a/vdslib/src/vespa/vdslib/container/searchresult.cpp
+++ b/vdslib/src/vespa/vdslib/container/searchresult.cpp
@@ -116,7 +116,8 @@ SearchResult::SearchResult(document::ByteBuffer & buf) :
     _numDocIdBytes(0),
     _aggregatorList(),
     _groupingList(),
-    _sortBlob()
+    _sortBlob(),
+    _match_features()
 {
     deserialize(buf);
 }
@@ -192,18 +193,18 @@ uint32_t SearchResult::getSerializedSize() const
 
 void SearchResult::addHit(uint32_t lid, const char * docId, RankType rank, const void * sortData, size_t sz)
 {
-    addHit(lid, docId, rank, _sortBlob.getCount());
+    addHit(lid, docId, rank);
     _sortBlob.append(sortData, sz);
 }
 
-void SearchResult::addHit(uint32_t lid, const char * docId, RankType rank, size_t index)
+void SearchResult::addHit(uint32_t lid, const char * docId, RankType rank)
 {
     const size_t sz(strlen(docId));
     size_t start = 0;
     if ( ! _hits.empty() ) {
         start = getBufCount();
     }
-    Hit h(lid, rank, start, index);
+    Hit h(lid, rank, start, _hits.size());
     _hits.push_back(h);
     _totalHits++;
     _numDocIdBytes += sz + 1;
@@ -222,4 +223,10 @@ void SearchResult::sort()
     }
 }
 
+void
+SearchResult::set_match_features(FeatureValues&& match_features)
+{
+    _match_features = std::move(match_features);
+}
+
 }
diff --git a/vdslib/src/vespa/vdslib/container/searchresult.h b/vdslib/src/vespa/vdslib/container/searchresult.h
index 2be2e5ee1c6..bc25ad76631 100644
--- a/vdslib/src/vespa/vdslib/container/searchresult.h
+++ b/vdslib/src/vespa/vdslib/container/searchresult.h
@@ -1,6 +1,8 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 #pragma once
 
+#include <vespa/vespalib/util/arrayref.h>
+#include <vespa/vespalib/util/featureset.h>
 #include <vespa/vespalib/util/memory.h>
 #include <vector>
 #include <map>
@@ -43,6 +45,8 @@ private:
 class SearchResult {
 public:
     using RankType = double;
+    using FeatureValue = vespalib::FeatureSet::Value;
+    using FeatureValues = vespalib::FeatureValues;
 public:
     SearchResult();
 
@@ -60,6 +64,11 @@ public:
     AggregatorList       & getAggregatorList()             { return _aggregatorList; }
     const AggregatorList & getAggregatorList()       const { return _aggregatorList; }
     void getSortBlob(size_t index, const void * & blob, size_t & sz) const { _sortBlob.getBlob(_hits[index].getIndex(), blob, sz); }
+    const FeatureValues& get_match_features() const noexcept { return _match_features; }
+    vespalib::ConstArrayRef<FeatureValue> get_match_feature_values(uint32_t index) const noexcept {
+        uint32_t num_features = _match_features.names.size();
+        return { _match_features.values.data() + _hits[index].getIndex() * num_features, num_features };
+    }
     void                   setWantedHitCount(size_t s)     { _wantedHits = s; }
     size_t                 getWantedHitCount()       const { return _wantedHits; }
     void                   setTotalHitCount(uint64_t s)    { _totalHits = s; }
@@ -73,8 +82,9 @@ public:
     void setRank(size_t hitNo, RankType rank) {
         _hits[hitNo].setRank(rank);
     }
-    void addHit(uint32_t lid, const char * docId, RankType rank, size_t index=0);
+    void addHit(uint32_t lid, const char * docId, RankType rank);
     void addHit(uint32_t lid, const char * docId, RankType rank, const void * sortData, size_t sz);
+    void set_match_features(FeatureValues&& match_features);
     void sort();
 
     void deserialize(document::ByteBuffer & buf);
@@ -98,7 +108,7 @@ private:
         uint32_t _lid;
         RankType _rank;
         uint32_t _docIdOffset;
-        uint32_t _index;  // refers to sortBlob
+        uint32_t _index;  // refers to _match_features and _sortBlob
     };
     class SortDataCompare {
     private:
@@ -131,6 +141,7 @@ private:
     AggregatorList               _aggregatorList;
     AggregatorList               _groupingList;
     BlobContainer                _sortBlob;
+    FeatureValues                _match_features;
 };
 
 }