change matchFeatures in LeanHit to be FeatureData directly

2 files changed, 16 insertions, 15 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 b948eab266a..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,11 +1,10 @@
 // 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.tensor.Tensor;
+import com.yahoo.data.access.Inspector;
+import com.yahoo.search.result.FeatureData;
 
 import java.util.Arrays;
-import java.util.Map;
-import java.util.HashMap;
 
 /**
  * @author baldersheim
@@ -17,7 +16,7 @@ public class LeanHit implements Comparable<LeanHit> {
     private final byte [] sortData;
     private final int partId;
     private final int distributionKey;
-    private final Map<String, Tensor> matchFeatures;
+    private FeatureData matchFeatures;
 
     public LeanHit(byte [] gid, int partId, int distributionKey, double relevance) {
         this(gid, partId, distributionKey, relevance, null);
@@ -28,7 +27,7 @@ public class LeanHit implements Comparable<LeanHit> {
         this.sortData = sortData;
         this.partId = partId;
         this.distributionKey = distributionKey;
-        this.matchFeatures = new HashMap<>();
+        this.matchFeatures = null;
     }
 
     public double getRelevance() { return relevance; }
@@ -37,10 +36,10 @@ public class LeanHit implements Comparable<LeanHit> {
     public boolean hasSortData() { return sortData != null; }
     public int getPartId() { return partId; }
     public int getDistributionKey() { return distributionKey; }
-    public final Map<String, Tensor> getMatchFeatures() { return matchFeatures; }
+    public FeatureData getMatchFeatures() { return matchFeatures; }
 
-    public void addMatchFeature(String name, Tensor value) {
-        matchFeatures.put(name, value);
+    public void addMatchFeatures(Inspector features) {
+        matchFeatures = new FeatureData(features);
     }
 
     @Override
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 98172d4a96e..54eeccd6505 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
@@ -230,19 +230,21 @@ public class ProtobufSerialization {
                     ? 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.getFeaturesList();
+                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.getTensorBlob();
-                    Tensor tensor = tensorBlob.isEmpty()
-                        ? Tensor.from(value.getValue())
-                        : TypedBinaryFormat.decode(Optional.empty(),
-                                                   GrowableByteBuffer.wrap(tensorBlob.toByteArray()));
-                    hit.addMatchFeature(name, tensor);
+                    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);
         }