Merge pull request #19779 from vespa-engine/arnej/add-match-features-in-protocol

add skeleton for match features

3 files changed, 40 insertions, 1 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/LeanHit.java b/container-search/src/main/java/com/yahoo/search/dispatch/LeanHit.java
index ae9d62fca41..6a47e19e310 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/LeanHit.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/LeanHit.java
@@ -1,6 +1,8 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.search.dispatch;
 
+import com.yahoo.data.access.Inspector;
+import com.yahoo.search.result.FeatureData;
 
 import java.util.Arrays;
 
@@ -14,6 +16,7 @@ public class LeanHit implements Comparable<LeanHit> {
     private final byte [] sortData;
     private final int partId;
     private final int distributionKey;
+    private FeatureData matchFeatures;
 
     public LeanHit(byte [] gid, int partId, int distributionKey, double relevance) {
         this(gid, partId, distributionKey, relevance, null);
@@ -24,6 +27,7 @@ public class LeanHit implements Comparable<LeanHit> {
         this.sortData = sortData;
         this.partId = partId;
         this.distributionKey = distributionKey;
+        this.matchFeatures = null;
     }
 
     public double getRelevance() { return relevance; }
@@ -32,6 +36,11 @@ public class LeanHit implements Comparable<LeanHit> {
     public boolean hasSortData() { return sortData != null; }
     public int getPartId() { return partId; }
     public int getDistributionKey() { return distributionKey; }
+    public FeatureData getMatchFeatures() { return matchFeatures; }
+
+    public void addMatchFeatures(Inspector features) {
+        matchFeatures = new FeatureData(features);
+    }
 
     @Override
     public int compareTo(LeanHit o) {
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
index a57fa6164d7..ac41321f639 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
@@ -33,6 +33,7 @@ import com.yahoo.vespa.objects.BufferSerializer;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.function.Consumer;
 
 public class ProtobufSerialization {
@@ -204,6 +205,7 @@ public class ProtobufSerialization {
         result.getResult().setCoverage(convertToCoverage(protobuf));
 
         convertSearchReplyErrors(result.getResult(), protobuf.getErrorsList());
+        List<String> featureNames = protobuf.getMatchFeatureNamesList();
 
         var haveGrouping = ! protobuf.getGroupingBlob().isEmpty();
         if (haveGrouping) {
@@ -224,6 +226,23 @@ public class ProtobufSerialization {
             LeanHit hit = (replyHit.getSortData().isEmpty())
                     ? new LeanHit(replyHit.getGlobalId().toByteArray(), partId, distKey, replyHit.getRelevance())
                     : new LeanHit(replyHit.getGlobalId().toByteArray(), partId, distKey, replyHit.getRelevance(), replyHit.getSortData().toByteArray());
+            if (! featureNames.isEmpty()) {
+                List<SearchProtocol.Feature> featureValues = replyHit.getMatchFeaturesList();
+                var object = new Value.ObjectValue();
+                var nameIter = featureNames.iterator();
+                var valueIter = featureValues.iterator();
+                while (nameIter.hasNext() && valueIter.hasNext()) {
+                    String name = nameIter.next();
+                    SearchProtocol.Feature value = valueIter.next();
+                    ByteString tensorBlob = value.getTensor();
+                    if (tensorBlob.isEmpty()) {
+                        object.put(name, value.getNumber());
+                    } else {
+                        object.put(name, new Value.DataValue(tensorBlob.toByteArray()));
+                    }
+                }
+                hit.addMatchFeatures(object);
+            }
             result.getLeanHits().add(hit);
         }
 
diff --git a/container-search/src/test/java/com/yahoo/search/handler/SearchHandlerTest.java b/container-search/src/test/java/com/yahoo/search/handler/SearchHandlerTest.java
index ffe5413c441..2c76383bff4 100644
--- a/container-search/src/test/java/com/yahoo/search/handler/SearchHandlerTest.java
+++ b/container-search/src/test/java/com/yahoo/search/handler/SearchHandlerTest.java
@@ -304,11 +304,22 @@ public class SearchHandlerTest {
         assertOkResult(driver.sendRequest(request), jsonResult);
     }
 
+    private boolean waitForMetric(String key) {
+        try {
+            for (int i = 0; i < 10; i++) {
+                if (metric.metrics().containsKey(key)) return true;
+                Thread.sleep(20);
+            }
+        } catch (InterruptedException e) {
+        }
+        return false;
+    }
+
     private void assertOkResult(RequestHandlerTestDriver.MockResponseHandler response, String expected) {
         assertEquals(expected, response.readAll());
         assertEquals(200, response.getStatus());
         assertEquals(selfHostname, response.getResponse().headers().get(myHostnameHeader).get(0));
-        assertTrue(metric.metrics().containsKey(SearchHandler.RENDER_LATENCY_METRIC));
+        assertTrue(waitForMetric(SearchHandler.RENDER_LATENCY_METRIC));
     }
 
     @Test