Add tensor short form format option for result rendering

6 files changed, 101 insertions, 36 deletions

diff --git a/container-search/abi-spec.json b/container-search/abi-spec.json
index 40071f90c34..b7ce40f19a2 100644
--- a/container-search/abi-spec.json
+++ b/container-search/abi-spec.json
@@ -7074,6 +7074,7 @@
     ],
     "methods": [
       "public void <init>(com.fasterxml.jackson.core.JsonGenerator, boolean)",
+      "public void <init>(com.fasterxml.jackson.core.JsonGenerator, boolean, boolean)",
       "public void accept(java.lang.String, java.lang.Object)",
       "public void accept(java.lang.String, byte[], int, int)",
       "protected boolean shouldRender(java.lang.String, java.lang.Object)",
@@ -7439,6 +7440,7 @@
       "public static com.yahoo.search.result.FeatureData empty()",
       "public com.yahoo.data.access.Inspector inspect()",
       "public java.lang.String toJson()",
+      "public java.lang.String toJson(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)",
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 0a87ad7ec2b..ee0e7f4fe0e 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
@@ -45,6 +45,7 @@ import com.yahoo.search.result.Hit;
 import com.yahoo.search.result.HitGroup;
 import com.yahoo.search.result.NanNumber;
 import com.yahoo.tensor.Tensor;
+import com.yahoo.tensor.serialization.JsonFormat;
 
 import java.io.IOException;
 import java.io.OutputStream;
@@ -58,7 +59,6 @@ import java.util.ArrayDeque;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Deque;
-import java.util.Iterator;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
@@ -78,6 +78,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
 
     private static final CompoundName DEBUG_RENDERING_KEY = new CompoundName("renderer.json.debug");
     private static final CompoundName JSON_CALLBACK = new CompoundName("jsoncallback");
+    private static final CompoundName TENSOR_FORMAT = new CompoundName("format.tensors");
 
     // if this must be optimized, simply use com.fasterxml.jackson.core.SerializableString
     private static final String BUCKET_LIMITS = "limits";
@@ -127,6 +128,8 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
     private LongSupplier timeSource;
     private OutputStream stream;
 
+    private boolean tensorShortFormRendering = false;
+
     public JsonRenderer() {
         this(null);
     }
@@ -166,6 +169,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
     public void beginResponse(OutputStream stream) throws IOException {
         beginJsonCallback(stream);
         debugRendering = getDebugRendering(getResult().getQuery());
+        tensorShortFormRendering = getTensorShortFormRendering(getResult().getQuery());
         setGenerator(generatorFactory.createGenerator(stream, JsonEncoding.UTF8), debugRendering);
         renderedChildren = new ArrayDeque<>();
         generator.writeStartObject();
@@ -200,6 +204,12 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
         return q != null && q.properties().getBoolean(DEBUG_RENDERING_KEY, false);
     }
 
+    private boolean getTensorShortFormRendering(Query q) {
+        if (q == null || q.properties().get(TENSOR_FORMAT) == null)
+            return false;
+        return q.properties().getString(TENSOR_FORMAT).equalsIgnoreCase("short");
+    }
+
     protected void renderTrace(Trace trace) throws IOException {
         if (!trace.traceNode().children().iterator().hasNext()) return;
         if (getResult().getQuery().getTraceLevel() == 0) return;
@@ -285,8 +295,6 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
             generator.writeEndObject();
         }
         generator.writeEndArray();
-
-
     }
 
     protected void renderCoverage() throws IOException {
@@ -510,7 +518,7 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
     }
 
     protected FieldConsumer createFieldConsumer(JsonGenerator generator, boolean debugRendering) {
-        return new FieldConsumer(generator, debugRendering);
+        return new FieldConsumer(generator, debugRendering, tensorShortFormRendering);
     }
 
     /**
@@ -529,12 +537,18 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
 
         private final JsonGenerator generator;
         private final boolean debugRendering;
+        private final boolean tensorShortForm;
 
         private MutableBoolean hasFieldsField;
 
         public FieldConsumer(JsonGenerator generator, boolean debugRendering) {
+            this(generator, debugRendering, false);
+        }
+
+        public FieldConsumer(JsonGenerator generator, boolean debugRendering, boolean tensorShortForm) {
             this.generator = generator;
             this.debugRendering = debugRendering;
+            this.tensorShortForm = tensorShortForm;
         }
 
         /**
@@ -659,7 +673,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());
+                generator.writeRawValue(((FeatureData)field).toJson(tensorShortForm));
             } else if (field instanceof Inspectable) {
                 renderInspectorDirect(((Inspectable)field).inspect());
             } else if (field instanceof JsonProducer) {
@@ -697,26 +711,18 @@ public class JsonRenderer extends AsynchronousSectionedRenderer<Result> {
         }
 
         private void renderTensor(Optional<Tensor> tensor) throws IOException {
-            generator.writeStartObject();
-            generator.writeArrayFieldStart("cells");
-            if (tensor.isPresent()) {
-                for (Iterator<Tensor.Cell> i = tensor.get().cellIterator(); i.hasNext(); ) {
-                    Tensor.Cell cell = i.next();
-
-                    generator.writeStartObject();
-
-                    generator.writeObjectFieldStart("address");
-                    for (int d = 0; d < cell.getKey().size(); d++)
-                        generator.writeObjectField(tensor.get().type().dimensions().get(d).name(), cell.getKey().label(d));
-                    generator.writeEndObject();
-
-                    generator.writeObjectField("value", cell.getValue());
-
-                    generator.writeEndObject();
-                }
+            if (tensor.isEmpty()) {
+                generator.writeStartObject();
+                generator.writeArrayFieldStart("cells");
+                generator.writeEndArray();
+                generator.writeEndObject();
+                return;
+            }
+            if (tensorShortForm) {
+                generator.writeRawValue(new String(JsonFormat.encodeShortForm(tensor.get()), StandardCharsets.UTF_8));
+            } else {
+                generator.writeRawValue(new String(JsonFormat.encode(tensor.get()), StandardCharsets.UTF_8));
             }
-            generator.writeEndArray();
-            generator.writeEndObject();
         }
 
     }
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 fd41d4ee10c..7673352576d 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
@@ -62,9 +62,17 @@ public class FeatureData implements Inspectable, JsonProducer {
         return jsonForm;
     }
 
+    public String toJson(boolean tensorShortForm) {
+        if (this == empty) return "{}";
+        if (jsonForm != null) return jsonForm;
+
+        jsonForm = JsonRender.render(value, new Encoder(new StringBuilder(), true, tensorShortForm)).toString();
+        return jsonForm;
+    }
+
     @Override
     public StringBuilder writeJson(StringBuilder target) {
-        return JsonRender.render(value, new Encoder(target, true));
+        return JsonRender.render(value, new Encoder(target, true, false));
     }
 
     /**
@@ -162,15 +170,19 @@ public class FeatureData implements Inspectable, JsonProducer {
     /** A JSON encoder which encodes DATA as a tensor */
     private static class Encoder extends JsonRender.StringEncoder {
 
-        Encoder(StringBuilder out, boolean compact) {
+        private final boolean tensorShortForm;
+
+        Encoder(StringBuilder out, boolean compact, boolean tensorShortForm) {
             super(out, compact);
+            this.tensorShortForm = tensorShortForm;
         }
 
         @Override
         public void encodeDATA(byte[] value) {
             // This could be done more efficiently ...
-            target().append(new String(JsonFormat.encodeWithType(TypedBinaryFormat.decode(Optional.empty(), GrowableByteBuffer.wrap(value))),
-                                       StandardCharsets.UTF_8));
+            Tensor tensor = TypedBinaryFormat.decode(Optional.empty(), GrowableByteBuffer.wrap(value));
+            byte[] encodedTensor = tensorShortForm ? JsonFormat.encodeShortForm(tensor) : JsonFormat.encodeWithType(tensor);
+            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 d0df131ae31..290b7266a3a 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
@@ -124,6 +124,55 @@ public class JsonRendererTestCase {
     }
 
     @Test
+    public void testTensorShortForm() throws ExecutionException, InterruptedException, IOException {
+        String expected = "{" +
+                "\"root\":{" +
+                    "\"id\":\"toplevel\"," +
+                    "\"relevance\":1.0," +
+                    "\"fields\":{" +
+                        "\"totalCount\":1" +
+                    "}," +
+                    "\"children\":[{" +
+                        "\"id\":\"tensors\"," +
+                        "\"relevance\":1.0," +
+                        "\"fields\":{" +
+                            "\"tensor_standard\":{\"type\":\"tensor(x{},y{})\",\"cells\":[{\"address\":{\"x\":\"a\",\"y\":\"0\"},\"value\":1.0},{\"address\":{\"x\":\"b\",\"y\":\"1\"},\"value\":2.0}]}," +
+                            "\"tensor_indexed\":{\"type\":\"tensor(x[2],y[3])\",\"values\":[[1.0,2.0,3.0],[4.0,5.0,6.0]]}," +
+                            "\"tensor_single_mapped\":{\"type\":\"tensor(x{})\",\"cells\":{\"a\":1.0,\"b\":2.0}}," +
+                            "\"tensor_mixed\":{\"type\":\"tensor(x{},y[2])\",\"blocks\":{\"a\":[1.0,2.0],\"b\":[3.0,4.0]}}," +
+                            "\"summaryfeatures\":{" +
+                                "\"tensor_standard\":{\"type\":\"tensor(x{},y{})\",\"cells\":[{\"address\":{\"x\":\"a\",\"y\":\"0\"},\"value\":1.0},{\"address\":{\"x\":\"b\",\"y\":\"1\"},\"value\":2.0}]}," +
+                                "\"tensor_indexed\":{\"type\":\"tensor(x[2],y[3])\",\"values\":[[1.0,2.0,3.0],[4.0,5.0,6.0]]}," +
+                                "\"tensor_single_mapped\":{\"type\":\"tensor(x{})\",\"cells\":{\"a\":1.0,\"b\":2.0}}," +
+                                "\"tensor_mixed\":{\"type\":\"tensor(x{},y[2])\",\"blocks\":{\"a\":[1.0,2.0],\"b\":[3.0,4.0]}}" +
+                            "}" +
+                        "}" +
+                    "}]" +
+                "}}\n";
+
+        Slime slime = new Slime();
+        Cursor features = slime.setObject();
+        features.setData("tensor_standard", TypedBinaryFormat.encode(Tensor.from("tensor(x{},y{}):{ {x:a,y:0}:1.0, {x:b,y:1}:2.0 }")));
+        features.setData("tensor_indexed", TypedBinaryFormat.encode(Tensor.from("tensor(x[2],y[3]):[[1,2,3],[4,5,6]]")));
+        features.setData("tensor_single_mapped", TypedBinaryFormat.encode(Tensor.from("tensor(x{}):{ a:1, b:2 }")));
+        features.setData("tensor_mixed", TypedBinaryFormat.encode(Tensor.from("tensor(x{},y[2]):{a:[1,2], b:[3,4]}")));
+        FeatureData summaryFeatures = new FeatureData(new SlimeAdapter(slime.get()));
+
+        Hit h = new Hit("tensors");
+        h.setField("tensor_standard", new TensorFieldValue(Tensor.from("tensor(x{},y{}):{ {x:a,y:0}:1.0, {x:b,y:1}:2.0 }")));
+        h.setField("tensor_indexed", new TensorFieldValue(Tensor.from("tensor(x[2],y[3]):[[1,2,3],[4,5,6]]")));
+        h.setField("tensor_single_mapped", new TensorFieldValue(Tensor.from("tensor(x{}):{ a:1, b:2 }")));
+        h.setField("tensor_mixed", new TensorFieldValue(Tensor.from("tensor(x{},y[2]):{a:[1,2], b:[3,4]}")));
+        h.setField("summaryfeatures", summaryFeatures);
+
+        Result r = new Result(new Query("/?format.tensors=short"));
+        r.hits().add(h);
+        r.setTotalHitCount(1L);
+        String summary = render(r);
+        assertEqualJson(expected, summary);
+    }
+
+    @Test
     public void testDataTypes() throws IOException, InterruptedException, ExecutionException {
         String expected = "{"
                 + "    \"root\": {"
@@ -463,7 +512,6 @@ public class JsonRendererTestCase {
         assertEqualJson(expected, summary);
     }
 
-
     @Test
     public void testTracingOfNodesWithBothChildrenAndDataAndEmptySubnode() throws IOException, InterruptedException, ExecutionException {
         String expected = "{"
@@ -553,7 +601,6 @@ public class JsonRendererTestCase {
         assertEqualJson(expected, summary);
     }
 
-
     @Test
     public void test() throws IOException, InterruptedException, ExecutionException {
         String expected = "{"
diff --git a/model-evaluation/src/main/java/ai/vespa/models/handler/ModelsEvaluationHandler.java b/model-evaluation/src/main/java/ai/vespa/models/handler/ModelsEvaluationHandler.java
index 9e365056355..8283d79fa2e 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/handler/ModelsEvaluationHandler.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/handler/ModelsEvaluationHandler.java
@@ -10,7 +10,6 @@ import com.yahoo.container.jdisc.ThreadedHttpRequestHandler;
 import com.yahoo.searchlib.rankingexpression.ExpressionFunction;
 import com.yahoo.slime.Cursor;
 import com.yahoo.slime.Slime;
-import com.yahoo.tensor.IndexedTensor;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorType;
 import com.yahoo.tensor.serialization.JsonFormat;
@@ -90,7 +89,7 @@ public class ModelsEvaluationHandler extends ThreadedHttpRequestHandler {
         }
         Tensor result = evaluator.evaluate();
 
-        Optional<String> format = property(request, "format");
+        Optional<String> format = property(request, "format.tensors");
         if (format.isPresent() && format.get().equalsIgnoreCase("short")) {
             return new Response(200, JsonFormat.encodeShortForm(result));
         }
diff --git a/model-evaluation/src/test/java/ai/vespa/models/handler/ModelsEvaluationHandlerTest.java b/model-evaluation/src/test/java/ai/vespa/models/handler/ModelsEvaluationHandlerTest.java
index 7029be24a60..992747b7866 100644
--- a/model-evaluation/src/test/java/ai/vespa/models/handler/ModelsEvaluationHandlerTest.java
+++ b/model-evaluation/src/test/java/ai/vespa/models/handler/ModelsEvaluationHandlerTest.java
@@ -1,7 +1,6 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.models.handler;
 
-import ai.vespa.models.evaluation.ModelTester;
 import ai.vespa.models.evaluation.ModelsEvaluator;
 import ai.vespa.models.evaluation.RankProfilesConfigImporterWithMockedConstants;
 import com.yahoo.config.subscription.ConfigGetter;
@@ -186,7 +185,7 @@ public class ModelsEvaluationHandlerTest {
     public void testMnistSoftmaxEvaluateSpecificFunctionWithShortOutput() {
         Map<String, String> properties = new HashMap<>();
         properties.put("Placeholder", inputTensorShortForm());
-        properties.put("format", "short");
+        properties.put("format.tensors", "short");
         String url = "http://localhost/model-evaluation/v1/mnist_softmax/default.add/eval";
         String expected = "{\"type\":\"tensor(d0[],d1[10])\",\"values\":[[-0.3546536862850189,0.3759574592113495,0.06054411828517914,-0.251544713973999,0.017951013520359993,1.2899067401885986,-0.10389615595340729,0.6367976665496826,-1.4136744737625122,-0.2573896050453186]]}";
         handler.assertResponse(url, properties, 200, expected);
@@ -216,7 +215,7 @@ public class ModelsEvaluationHandlerTest {
     @Test
     public void testVespaModelShortOutput() {
         Map<String, String> properties = new HashMap<>();
-        properties.put("format", "short");
+        properties.put("format.tensors", "short");
         String url = "http://localhost/model-evaluation/v1/vespa_model/";
         handler.assertResponse(url + "test_mapped/eval", properties, 200,
                 "{\"type\":\"tensor(d0{})\",\"cells\":{\"a\":1.0,\"b\":2.0}}");
@@ -231,7 +230,7 @@ public class ModelsEvaluationHandlerTest {
     @Test
     public void testVespaModelLiteralOutput() {
         Map<String, String> properties = new HashMap<>();
-        properties.put("format", "string");
+        properties.put("format.tensors", "string");
         String url = "http://localhost/model-evaluation/v1/vespa_model/";
         handler.assertResponse(url + "test_mapped/eval", properties, 200,
                 "tensor(d0{}):{a:1.0,b:2.0}");