Merge pull request #29004 from vespa-engine/arnej/add-get-as-tensor-with-defaults

add getAsTensor() API in RankProperties

1 files changed, 22 insertions, 0 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/query/ranking/RankProperties.java b/container-search/src/main/java/com/yahoo/search/query/ranking/RankProperties.java
index 544f26a7d89..4ac5375807b 100644
--- a/container-search/src/main/java/com/yahoo/search/query/ranking/RankProperties.java
+++ b/container-search/src/main/java/com/yahoo/search/query/ranking/RankProperties.java
@@ -3,6 +3,7 @@ package com.yahoo.search.query.ranking;
 
 import com.yahoo.fs4.GetDocSumsPacket;
 import com.yahoo.fs4.MapEncoder;
+import com.yahoo.tensor.Tensor;
 import com.yahoo.text.JSON;
 
 import java.nio.ByteBuffer;
@@ -11,6 +12,7 @@ import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 /**
  * Contains the properties of a query.
@@ -61,6 +63,26 @@ public class RankProperties implements Cloneable {
         return Collections.unmodifiableList(stringValues);
     }
 
+    /**
+     * Returns a tensor (as moved from RankFeatures by prepare step) if present
+     *
+     * @throws IllegalArgumentException if the value is there but wrong type
+     */
+    public Optional<Tensor> getAsTensor(String name) {
+        List<Object> values = properties.get(name);
+        if (values == null || values.isEmpty()) return Optional.empty();
+        if (values.size() != 1) {
+            throw new IllegalArgumentException("unexpected multiple [" + values.size() + "] values for property '" + name + "'");
+        }
+        Object feature = values.get(0);
+        if (feature == null) return Optional.empty();
+        if (feature instanceof Tensor t) return Optional.of(t);
+        if (feature instanceof Double d) return Optional.of(Tensor.from(d));
+        throw new IllegalArgumentException("Expected '" + name + "' to be a tensor or double, but it is '" + feature +
+                "', this usually means that '" + name + "' is not defined in the schema. " +
+                "See https://docs.vespa.ai/en/tensor-user-guide.html#querying-with-tensors");
+    }
+
     /** Removes all properties for a given name */
     public void remove(String name) {
         properties.remove(name);