Support direct tensor rendering

6 files changed, 71 insertions, 31 deletions

diff --git a/container-search/abi-spec.json b/container-search/abi-spec.json
index b70a88d09a0..a71082ecba3 100644
--- a/container-search/abi-spec.json
+++ b/container-search/abi-spec.json
@@ -899,19 +899,19 @@
     ],
     "fields" : [ ]
   },
-  "com.yahoo.prelude.query.MultiRangeItem$Limit": {
-    "superClass": "java.lang.Enum",
-    "interfaces": [],
-    "attributes": [
+  "com.yahoo.prelude.query.MultiRangeItem$Limit" : {
+    "superClass" : "java.lang.Enum",
+    "interfaces" : [ ],
+    "attributes" : [
       "public",
       "final",
       "enum"
     ],
-    "methods": [
+    "methods" : [
       "public static com.yahoo.prelude.query.MultiRangeItem$Limit[] values()",
       "public static com.yahoo.prelude.query.MultiRangeItem$Limit valueOf(java.lang.String)"
     ],
-    "fields": [
+    "fields" : [
       "public static final enum com.yahoo.prelude.query.MultiRangeItem$Limit INCLUSIVE",
       "public static final enum com.yahoo.prelude.query.MultiRangeItem$Limit EXCLUSIVE"
     ]
@@ -5318,7 +5318,10 @@
       "public void setSummaryFields(java.lang.String)",
       "public boolean getTensorShortForm()",
       "public void setTensorShortForm(java.lang.String)",
+      "public void setTensorFormat(java.lang.String)",
       "public void setTensorShortForm(boolean)",
+      "public boolean getTensorDirectValues()",
+      "public void setTensorDirectValues(boolean)",
       "public void prepare()",
       "public boolean equals(java.lang.Object)",
       "public int hashCode()"
@@ -7678,6 +7681,7 @@
       "public com.yahoo.data.access.Inspector inspect()",
       "public java.lang.String toJson()",
       "public java.lang.String toJson(boolean)",
+      "public java.lang.String toJson(boolean, boolean)",
       "public java.lang.StringBuilder writeJson(java.lang.StringBuilder)",
       "public java.lang.Double getDouble(java.lang.String)",
       "public com.yahoo.tensor.Tensor getTensor(java.lang.String)",
@@ -8725,4 +8729,4 @@
     ],
     "fields" : [ ]
   }
-}
+}
\ No newline at end of file
diff --git a/container-search/src/main/java/com/yahoo/search/query/Presentation.java b/container-search/src/main/java/com/yahoo/search/query/Presentation.java
index fd4519fdbb0..b949d1edabd 100644
--- a/container-search/src/main/java/com/yahoo/search/query/Presentation.java
+++ b/container-search/src/main/java/com/yahoo/search/query/Presentation.java
@@ -77,6 +77,9 @@ public class Presentation implements Cloneable {
     /** Whether to renders tensors in short form */
     private boolean tensorShortForm = true;
 
+    /** Whether to renders tensors in short form */
+    private boolean tensorDirectValues = false; // TODO: Flip default on Vespa 9
+
     /** Set of explicitly requested summary fields, instead of summary classes */
     private Set<String> summaryFields = LazySet.newHashSet();
 
@@ -178,27 +181,41 @@ public class Presentation implements Cloneable {
 
     /**
      * Returns whether tensors should use short form in JSON and textual representations, see
-     * <a href="https://docs.vespa.ai/en/reference/document-json-format.html#tensor">https://docs.vespa.ai/en/reference/document-json-format.html#tensor</a>
-     * and <a href="https://docs.vespa.ai/en/reference/tensor.html#tensor-literal-form">https://docs.vespa.ai/en/reference/tensor.html#tensor-literal-form</a>.
+     * <a href="https://docs.vespa.ai/en/reference/document-json-format.html#tensor">https://docs.vespa.ai/en/reference/document-json-format.html#tensor</a>.
      * Default is true.
      */
     public boolean getTensorShortForm() { return tensorShortForm; }
 
+    /** @deprecated use setTensorFormat(). */
+    @Deprecated // TODO: Remove on Vespa 9
+    public void setTensorShortForm(String value) {
+        setTensorFormat(value);
+    }
     /**
      * Sets whether tensors should use short form in JSON and textual representations from a string.
      *
      * @param value a string which must be either 'short' or 'long'
      * @throws IllegalArgumentException if any other value is passed
      */
-    public void setTensorShortForm(String value) {
-        tensorShortForm = toTensorShortForm(value);
-    }
-
-    private boolean toTensorShortForm(String value) {
-        return switch (value) {
-            case "short" -> true;
-            case "long" -> false;
-            default -> throw new IllegalArgumentException("Value must be 'long' or 'short', not '" + value + "'");
+    public void setTensorFormat(String value) {
+        switch (value) {
+            case "short" :
+                tensorShortForm = true;
+                tensorDirectValues = false;
+                break;
+            case "long" :
+                tensorShortForm = false;
+                tensorDirectValues = false;
+                break;
+            case "short-value" :
+                tensorShortForm = true;
+                tensorDirectValues = true;
+                break;
+            case "long-value" :
+                tensorShortForm = false;
+                tensorDirectValues = true;
+                break;
+            default : throw new IllegalArgumentException("Value must be 'long', 'short', 'long-value', or 'short-value', not '" + value + "'");
         };
     }
 
@@ -206,6 +223,19 @@ public class Presentation implements Cloneable {
         this.tensorShortForm = tensorShortForm;
     }
 
+    /**
+     * Returns whether tensor content should be rendered directly, or inside a JSON object containing a
+     * "type" entry having the tensor type, and a "cells"/"values"/"blocks" entry (depending on type),
+     * having the tensor content. See
+     * <a href="https://docs.vespa.ai/en/reference/document-json-format.html#tensor">https://docs.vespa.ai/en/reference/document-json-format.html#tensor</a>.
+     * Default is false: Render wrapped in a JSON object.
+     */
+    public boolean getTensorDirectValues() { return tensorDirectValues; }
+
+    public void setTensorDirectValues(boolean tensorDirectValues) {
+        this.tensorDirectValues = tensorDirectValues;
+    }
+
     /** Prepares this for binary serialization. For internal use - see {@link Query#prepare} */
     public void prepare() {
         if (highlight != null)
diff --git a/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java b/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java
index d5dc8120f29..e4a83972fae 100644
--- a/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java
+++ b/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java
@@ -303,7 +303,7 @@ public class QueryProperties extends Properties {
                 }
                 else if (key.size() == 3 && key.get(1).equals(Presentation.FORMAT)) {
                     if (key.last().equals(Presentation.TENSORS))
-                        query.getPresentation().setTensorShortForm(asString(value, "short"));
+                        query.getPresentation().setTensorFormat(asString(value, "short")); // TODO: Switch default to short-value on Vespa 9
                     else
                         throwIllegalParameter(key.last(), Presentation.FORMAT);
                 }
diff --git a/container-search/src/main/java/com/yahoo/search/rendering/JsonRenderer.java b/container-search/src/main/java/com/yahoo/search/rendering/JsonRenderer.java
index 9bb7e882a4b..9498f860f88 100644
--- a/container-search/src/main/java/com/yahoo/search/rendering/JsonRenderer.java
+++ b/container-search/src/main/java/com/yahoo/search/rendering/JsonRenderer.java
@@ -132,6 +132,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
         volatile boolean jsonMapsAll = true;
         volatile boolean jsonWsetsAll = false;
         volatile boolean tensorShortForm = true;
+        volatile boolean tensorDirectValues = false;
         boolean convertDeep() { return (jsonDeepMaps || jsonWsets); }
         void init() {
             this.debugRendering = false;
@@ -140,6 +141,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
             this.jsonMapsAll = true;
             this.jsonWsetsAll = true;
             this.tensorShortForm = true;
+            this.tensorDirectValues = false;
         }
         void getSettings(Query q) {
             if (q == null) {
@@ -154,7 +156,8 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
             this.jsonMapsAll = props.getBoolean(WRAP_DEEP_MAPS, true);
             this.jsonWsetsAll = props.getBoolean(WRAP_WSETS, true);
             this.tensorShortForm = q.getPresentation().getTensorShortForm();
-        }
+            this.tensorDirectValues = q.getPresentation().getTensorDirectValues();
+            }
     }
 
     private volatile FieldConsumerSettings fieldConsumerSettings;
@@ -776,7 +779,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
             } else if (field instanceof Tensor) {
                 renderTensor(Optional.of((Tensor)field));
             } else if (field instanceof FeatureData) {
-                generator().writeRawValue(((FeatureData)field).toJson(settings.tensorShortForm));
+                generator().writeRawValue(((FeatureData)field).toJson(settings.tensorShortForm, settings.tensorDirectValues));
             } else if (field instanceof Inspectable) {
                 renderInspectorDirect(((Inspectable)field).inspect());
             } else if (field instanceof JsonProducer) {
@@ -821,11 +824,8 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
                 generator().writeEndObject();
                 return;
             }
-            if (settings.tensorShortForm && 1==2) {
-                generator().writeRawValue(new String(JsonFormat.encodeShortForm(tensor.get()), StandardCharsets.UTF_8));
-            } else {
-                generator().writeRawValue(new String(JsonFormat.encode(tensor.get()), StandardCharsets.UTF_8));
-            }
+            generator().writeRawValue(new String(JsonFormat.encode(tensor.get(), settings.tensorShortForm, settings.tensorDirectValues),
+                                                 StandardCharsets.UTF_8));
         }
 
         private JsonGenerator generator() {
diff --git a/container-search/src/main/java/com/yahoo/search/result/FeatureData.java b/container-search/src/main/java/com/yahoo/search/result/FeatureData.java
index 2cb5e0e07e9..421f19475a6 100644
--- a/container-search/src/main/java/com/yahoo/search/result/FeatureData.java
+++ b/container-search/src/main/java/com/yahoo/search/result/FeatureData.java
@@ -65,16 +65,20 @@ public class FeatureData implements Inspectable, JsonProducer {
     }
 
     public String toJson(boolean tensorShortForm) {
+        return toJson(tensorShortForm, false);
+    }
+
+    public String toJson(boolean tensorShortForm, boolean tensorDirectValues) {
         if (this == empty) return "{}";
         if (jsonForm != null) return jsonForm;
 
-        jsonForm = JsonRender.render(value, new Encoder(new StringBuilder(), true, tensorShortForm)).toString();
+        jsonForm = JsonRender.render(value, new Encoder(new StringBuilder(), true, tensorShortForm, tensorDirectValues)).toString();
         return jsonForm;
     }
 
     @Override
     public StringBuilder writeJson(StringBuilder target) {
-        return JsonRender.render(value, new Encoder(target, true, false));
+        return JsonRender.render(value, new Encoder(target, true, false, false));
     }
 
     /**
@@ -173,17 +177,19 @@ public class FeatureData implements Inspectable, JsonProducer {
     private static class Encoder extends JsonRender.StringEncoder {
 
         private final boolean tensorShortForm;
+        private final boolean tensorDirectValues;
 
-        Encoder(StringBuilder out, boolean compact, boolean tensorShortForm) {
+        Encoder(StringBuilder out, boolean compact, boolean tensorShortForm, boolean tensorDirectValues) {
             super(out, compact);
             this.tensorShortForm = tensorShortForm;
+            this.tensorDirectValues = tensorDirectValues;
         }
 
         @Override
         public void encodeDATA(byte[] value) {
             // This could be done more efficiently ...
             Tensor tensor = TypedBinaryFormat.decode(Optional.empty(), GrowableByteBuffer.wrap(value));
-            byte[] encodedTensor = tensorShortForm ? JsonFormat.encodeShortForm(tensor) : JsonFormat.encodeWithType(tensor);
+            byte[] encodedTensor = JsonFormat.encode(tensor, tensorShortForm, tensorDirectValues);
             target().append(new String(encodedTensor, StandardCharsets.UTF_8));
         }
 
diff --git a/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java b/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java
index 9486eeb92de..c1ede03a371 100644
--- a/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java
@@ -219,7 +219,7 @@ public class JsonRendererTestCase {
             fail("Expected exception");
         }
         catch (IllegalArgumentException e) {
-            assertEquals("Could not set 'presentation.format.tensors' to 'unknown': Value must be 'long' or 'short', not 'unknown'",
+            assertEquals("Could not set 'presentation.format.tensors' to 'unknown': Value must be 'long', 'short', 'long-value', or 'short-value', not 'unknown'",
                     Exceptions.toMessageString(e));
         }
     }