2 files changed, 32 insertions, 15 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);
diff --git a/container-search/src/main/java/com/yahoo/search/ranking/PreparedInput.java b/container-search/src/main/java/com/yahoo/search/ranking/PreparedInput.java
index 346acccd916..5491724cc08 100644
--- a/container-search/src/main/java/com/yahoo/search/ranking/PreparedInput.java
+++ b/container-search/src/main/java/com/yahoo/search/ranking/PreparedInput.java
@@ -26,24 +26,19 @@ record PreparedInput(String name, Tensor value) {
         List<PreparedInput> result = new ArrayList<>();
         var ranking = query.getRanking();
         var rankFeatures = ranking.getFeatures();
-        var rankProps = ranking.getProperties().asMap();
+        var rankProps = ranking.getProperties();
         for (String queryFeatureName : queryFeatures) {
             String needed = "query(" + queryFeatureName + ")";
-            // searchers are recommended to place query features here:
-            var feature = rankFeatures.getTensor(needed);
-            if (feature.isPresent()) {
-                result.add(new PreparedInput(needed, feature.get()));
-            } else {
-                // but other ways of setting query features end up in the properties:
-                var objList = rankProps.get(queryFeatureName);
-                if (objList != null && objList.size() == 1 && objList.get(0) instanceof Tensor t) {
-                    result.add(new PreparedInput(needed, t));
-                } else if (objList != null && objList.size() == 1 && objList.get(0) instanceof Double d) {
-                    result.add(new PreparedInput(needed, Tensor.from(d)));
-                } else {
-                    throw new IllegalArgumentException("missing query feature: " + queryFeatureName);
-                }
+            // after prepare() the query tensor ends up here:
+            var feature = rankProps.getAsTensor(queryFeatureName);
+            if (feature.isEmpty()) {
+                // searchers are recommended to place query features here:
+                feature = rankFeatures.getTensor(needed);
             }
+            if (feature.isEmpty()) {
+                throw new IllegalArgumentException("missing query feature: " + queryFeatureName);
+            }
+            result.add(new PreparedInput(needed, feature.get()));
         }
         return result;
     }