Merge pull request #9536 from vespa-engine/bratseth/double-api-in-rank-features

Add double API to RankFeatures

6 files changed, 67 insertions, 16 deletions

diff --git a/container-search/abi-spec.json b/container-search/abi-spec.json
index fb50da7bff1..e96d208679d 100644
--- a/container-search/abi-spec.json
+++ b/container-search/abi-spec.json
@@ -6425,10 +6425,12 @@
     ],
     "methods": [
       "public void <init>()",
-      "public void put(java.lang.String, java.lang.String)",
+      "public void put(java.lang.String, double)",
       "public void put(java.lang.String, com.yahoo.tensor.Tensor)",
+      "public void put(java.lang.String, java.lang.String)",
       "public java.lang.String get(java.lang.String)",
       "public java.lang.Object getObject(java.lang.String)",
+      "public java.util.OptionalDouble getDouble(java.lang.String)",
       "public java.util.Optional getTensor(java.lang.String)",
       "public java.util.Map asMap()",
       "public boolean isEmpty()",
diff --git a/container-search/src/main/java/com/yahoo/fs4/MapEncoder.java b/container-search/src/main/java/com/yahoo/fs4/MapEncoder.java
index 565a4c483c3..3e45d7e5f53 100644
--- a/container-search/src/main/java/com/yahoo/fs4/MapEncoder.java
+++ b/container-search/src/main/java/com/yahoo/fs4/MapEncoder.java
@@ -64,12 +64,12 @@ public class MapEncoder {
         buffer.putInt(utf8.length);
         buffer.put(utf8);
         buffer.putInt(map.size());
-        for (Map.Entry<String, ?> property : map.entrySet()) {
-            String key = property.getKey();
+        for (Map.Entry<String, ?> entry : map.entrySet()) {
+            String key = entry.getKey();
             utf8 = Utf8.toBytes(key);
             buffer.putInt(utf8.length);
             buffer.put(utf8);
-            Object value = property.getValue();
+            Object value = entry.getValue();
             if (value == null) {
                 utf8 = Utf8.toBytes("");
             } else {
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/MapConverter.java b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/MapConverter.java
index 74828dd6740..7c97da6a77d 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/MapConverter.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/MapConverter.java
@@ -16,6 +16,7 @@ import java.util.function.Consumer;
  * @author ollivir
  */
 public class MapConverter {
+
     public static void convertMapTensors(Map<String, Object> map, Consumer<TensorProperty.Builder> inserter) {
         for (var entry : map.entrySet()) {
             var value = entry.getValue();
@@ -26,7 +27,7 @@ public class MapConverter {
         }
     }
 
-    public static void convertMapStrings(Map<String, Object> map, Consumer<StringProperty.Builder> inserter) {
+    public static void convertMapPrimitives(Map<String, Object> map, Consumer<StringProperty.Builder> inserter) {
         for (var entry : map.entrySet()) {
             var value = entry.getValue();
             if (!(value instanceof Tensor)) {
@@ -66,4 +67,5 @@ public class MapConverter {
             }
         }
     }
+
 }
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 5ef4562f040..39a1587afea 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
@@ -31,6 +31,7 @@ import java.util.List;
 import java.util.function.Consumer;
 
 public class ProtobufSerialization {
+
     private static final int INITIAL_SERIALIZATION_BUFFER_SIZE = 10 * 1024;
 
     public static byte[] serializeSearchRequest(Query query, String serverId) {
@@ -87,7 +88,7 @@ public class ProtobufSerialization {
         }
 
         var featureMap = ranking.getFeatures().asMap();
-        MapConverter.convertMapStrings(featureMap, builder::addFeatureOverrides);
+        MapConverter.convertMapPrimitives(featureMap, builder::addFeatureOverrides);
         MapConverter.convertMapTensors(featureMap, builder::addTensorFeatureOverrides);
         mergeRankProperties(ranking, builder::addRankProperties, builder::addTensorRankProperties);
     }
@@ -101,8 +102,10 @@ public class ProtobufSerialization {
         }
     }
 
-    public static SearchProtocol.DocsumRequest.Builder createDocsumRequestBuilder(Query query, String serverId, String summaryClass,
-            boolean includeQueryData) {
+    public static SearchProtocol.DocsumRequest.Builder createDocsumRequestBuilder(Query query,
+                                                                                  String serverId,
+                                                                                  String summaryClass,
+                                                                                  boolean includeQueryData) {
         var builder = SearchProtocol.DocsumRequest.newBuilder()
                 .setTimeout((int) query.getTimeLeft())
                 .setDumpFeatures(query.properties().getBoolean(Ranking.RANKFEATURES, false));
@@ -146,7 +149,7 @@ public class ProtobufSerialization {
         if (ranking.getLocation() != null) {
             builder.setGeoLocation(ranking.getLocation().toString());
         }
-        MapConverter.convertMapStrings(featureMap, builder::addFeatureOverrides);
+        MapConverter.convertMapPrimitives(featureMap, builder::addFeatureOverrides);
         MapConverter.convertMapTensors(featureMap, builder::addTensorFeatureOverrides);
         if (query.getPresentation().getHighlight() != null) {
             MapConverter.convertStringMultiMap(query.getPresentation().getHighlight().getHighlightTerms(), builder::addHighlightTerms);
@@ -165,8 +168,12 @@ public class ProtobufSerialization {
         return result;
     }
 
-    private static Result convertToResult(Query query, SearchProtocol.SearchReply protobuf, DocumentDatabase documentDatabase, int partId,
-            int distKey, String source) {
+    private static Result convertToResult(Query query,
+                                          SearchProtocol.SearchReply protobuf,
+                                          DocumentDatabase documentDatabase,
+                                          int partId,
+                                          int distKey,
+                                          String source) {
         var result = new Result(query);
 
         result.setTotalHitCount(protobuf.getTotalHitCount());
@@ -269,12 +276,14 @@ public class ProtobufSerialization {
         }
     }
 
-    private static void mergeRankProperties(Ranking ranking, Consumer<StringProperty.Builder> stringProperties,
-            Consumer<TensorProperty.Builder> tensorProperties) {
+    private static void mergeRankProperties(Ranking ranking,
+                                            Consumer<StringProperty.Builder> stringProperties,
+                                            Consumer<TensorProperty.Builder> tensorProperties) {
         MapConverter.convertMultiMap(ranking.getProperties().asMap(), propB -> {
             if (!GetDocSumsPacket.sessionIdKey.equals(propB.getName())) {
                 stringProperties.accept(propB);
             }
         }, tensorProperties);
     }
+
 }
diff --git a/container-search/src/main/java/com/yahoo/search/query/ranking/RankFeatures.java b/container-search/src/main/java/com/yahoo/search/query/ranking/RankFeatures.java
index 9786eba163a..3bc3a629c5d 100644
--- a/container-search/src/main/java/com/yahoo/search/query/ranking/RankFeatures.java
+++ b/container-search/src/main/java/com/yahoo/search/query/ranking/RankFeatures.java
@@ -11,6 +11,7 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.OptionalDouble;
 
 /**
  * Contains the rank features of a query.
@@ -29,8 +30,8 @@ public class RankFeatures implements Cloneable {
         this.features = features;
     }
 
-    /** Sets a rank feature by full name to a value */
-    public void put(String name, String value) {
+    /** Sets a double rank feature */
+    public void put(String name, double value) {
         features.put(name, value);
     }
 
@@ -39,6 +40,11 @@ public class RankFeatures implements Cloneable {
         features.put(name, value);
     }
 
+    /** Sets a rank feature to a value represented as a string */
+    public void put(String name, String value) {
+        features.put(name, value);
+    }
+
     /** Returns a rank feature as a string by full name or null if not set */
     public String get(String name) {
         Object value = features.get(name);
@@ -52,6 +58,18 @@ public class RankFeatures implements Cloneable {
     }
 
     /**
+     * Returns a double rank feature, or empty if there is no value with this name.
+     *
+     * @throws IllegalArgumentException if the value is set but is not a double
+     */
+    public OptionalDouble getDouble(String name) {
+        Object feature = features.get(name);
+        if (feature == null) return OptionalDouble.empty();
+        if (feature instanceof Double) return OptionalDouble.of((Double)feature);
+        throw new IllegalArgumentException("Expected a double value of '" + name + "' but has " + feature);
+    }
+
+    /**
      * Returns a tensor rank feature, or empty if there is no value with this name.
      *
      * @throws IllegalArgumentException if the value is set but is not a tensor
diff --git a/container-search/src/test/java/com/yahoo/fs4/test/RankFeaturesTestCase.java b/container-search/src/test/java/com/yahoo/fs4/test/RankFeaturesTestCase.java
index e8c16e572ae..b52c708ce4b 100644
--- a/container-search/src/test/java/com/yahoo/fs4/test/RankFeaturesTestCase.java
+++ b/container-search/src/test/java/com/yahoo/fs4/test/RankFeaturesTestCase.java
@@ -24,12 +24,31 @@ public class RankFeaturesTestCase {
     public void requireThatRankPropertiesTakesBothStringAndObject() {
         RankProperties p = new RankProperties();
         p.put("string", "b");
-        p.put("object", Integer.valueOf(7));
+        p.put("object", 7);
         assertEquals("7", p.get("object").get(0));
         assertEquals("b", p.get("string").get(0));
     }
 
     @Test
+    public void requireThatRankFeaturesUsingDoubleAndDoubleToStringEncodeTheSameWay() {
+        RankFeatures withDouble = new RankFeatures();
+        withDouble.put("query(myDouble)", 3.8);
+        assertEquals(3.8, withDouble.getDouble("query(myDouble)").getAsDouble(), 0.000001);
+
+        RankFeatures withString = new RankFeatures();
+        withString.put("query(myDouble)", String.valueOf(3.8));
+
+        RankProperties withDoubleP = new RankProperties();
+        withDouble.prepare(withDoubleP);
+        RankProperties withStringP = new RankProperties();
+        withString.prepare(withStringP);
+
+        byte[] withDoubleEncoded = encode(withDoubleP);
+        byte[] withStringEncoded = encode(withStringP);
+        assertEquals(Arrays.toString(withStringEncoded), Arrays.toString(withDoubleEncoded));
+    }
+
+    @Test
     public void requireThatSingleTensorIsBinaryEncoded() {
         TensorType type = new TensorType.Builder().mapped("x").mapped("y").mapped("z").build();
         Tensor tensor = Tensor.from(type, "{ {x:a, y:b, z:c}:2.0, {x:a, y:b, z:c2}:3.0 }");
@@ -113,4 +132,5 @@ public class RankFeaturesTestCase {
         }
         return result;
     }
+
 }