Add tensor short form rendering to document/v1

3 files changed, 63 insertions, 2 deletions

diff --git a/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java b/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java
index 340bd542885..a9d34fc78d8 100644
--- a/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java
+++ b/document/src/main/java/com/yahoo/document/json/JsonSerializationHelper.java
@@ -32,10 +32,12 @@ import com.yahoo.document.serialization.FieldWriter;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorAddress;
 import com.yahoo.tensor.TensorType;
+import com.yahoo.tensor.serialization.JsonFormat;
 import com.yahoo.vespa.objects.FieldBase;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
 import java.util.Base64;
 import java.util.Iterator;
 import java.util.Map;
@@ -71,6 +73,19 @@ public class JsonSerializationHelper {
         }
     }
 
+    public static void serializeTensorFieldShortForm(JsonGenerator generator, FieldBase field, TensorFieldValue value) {
+        wrapIOException(() -> {
+            fieldNameIfNotNull(generator, field);
+            if (value.getTensor().isPresent()) {
+                Tensor tensor = value.getTensor().get();
+                generator.writeRawValue(new String(JsonFormat.encodeShortForm(tensor), StandardCharsets.UTF_8));
+            } else {
+                generator.writeStartObject();
+                generator.writeEndObject();
+            }
+        });
+    }
+
     public static void serializeTensorField(JsonGenerator generator, FieldBase field, TensorFieldValue value) {
         wrapIOException(() -> {
             fieldNameIfNotNull(generator, field);
diff --git a/document/src/main/java/com/yahoo/document/json/JsonWriter.java b/document/src/main/java/com/yahoo/document/json/JsonWriter.java
index 95faf753877..05147d2c0ab 100644
--- a/document/src/main/java/com/yahoo/document/json/JsonWriter.java
+++ b/document/src/main/java/com/yahoo/document/json/JsonWriter.java
@@ -63,6 +63,7 @@ import static com.yahoo.document.json.JsonSerializationHelper.serializeStringFie
 import static com.yahoo.document.json.JsonSerializationHelper.serializeStructField;
 import static com.yahoo.document.json.JsonSerializationHelper.serializeStructuredField;
 import static com.yahoo.document.json.JsonSerializationHelper.serializeTensorField;
+import static com.yahoo.document.json.JsonSerializationHelper.serializeTensorFieldShortForm;
 import static com.yahoo.document.json.JsonSerializationHelper.serializeWeightedSet;
 import static com.yahoo.document.json.document.DocumentParser.FIELDS;
 import static com.yahoo.document.json.document.DocumentParser.REMOVE;
@@ -77,6 +78,8 @@ public class JsonWriter implements DocumentWriter {
     private static final JsonFactory jsonFactory = new JsonFactory();
     private final JsonGenerator generator;
 
+    private final boolean tensorShortForm;
+
     // I really hate exception unsafe constructors, but the alternative
     // requires generator to not be a final
     /**
@@ -108,7 +111,12 @@ public class JsonWriter implements DocumentWriter {
      *            the output JSON generator
      */
     public JsonWriter(JsonGenerator generator) {
+        this(generator, false);
+    }
+
+    public JsonWriter(JsonGenerator generator, boolean tensorShortForm) {
         this.generator = generator;
+        this.tensorShortForm = tensorShortForm;
     }
 
     /**
@@ -206,7 +214,11 @@ public class JsonWriter implements DocumentWriter {
 
     @Override
     public void write(FieldBase field, TensorFieldValue value) {
-        serializeTensorField(generator, field, value);
+        if (tensorShortForm) {
+            serializeTensorFieldShortForm(generator, field, value);
+        } else {
+            serializeTensorField(generator, field, value);
+        }
     }
 
     @Override
diff --git a/document/src/test/java/com/yahoo/document/json/JsonWriterTestCase.java b/document/src/test/java/com/yahoo/document/json/JsonWriterTestCase.java
index 7573aba519f..b7f24368ad7 100644
--- a/document/src/test/java/com/yahoo/document/json/JsonWriterTestCase.java
+++ b/document/src/test/java/com/yahoo/document/json/JsonWriterTestCase.java
@@ -2,6 +2,7 @@
 package com.yahoo.document.json;
 
 import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.io.JsonStringEncoder;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.yahoo.document.ArrayDataType;
@@ -25,6 +26,7 @@ import com.yahoo.document.datatypes.TensorFieldValue;
 import com.yahoo.document.internal.GeoPosType;
 import com.yahoo.document.json.readers.DocumentParseInfo;
 import com.yahoo.document.json.readers.VespaJsonDocumentReader;
+import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorType;
 import com.yahoo.text.Utf8;
 import org.junit.After;
@@ -32,6 +34,7 @@ import org.junit.Before;
 import org.junit.Test;
 
 import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.HashMap;
@@ -73,6 +76,7 @@ public class JsonWriterTestCase {
         registerSinglePositionDocumentType();
         registerMultiPositionDocumentType();
         registerTensorDocumentType();
+        registerIndexedTensorDocumentType();
         registerReferenceDocumentType();
     }
 
@@ -90,6 +94,13 @@ public class JsonWriterTestCase {
         types.registerDocumentType(x);
     }
 
+    private void registerIndexedTensorDocumentType() {
+        DocumentType x = new DocumentType("testindexedtensor");
+        TensorType tensorType = new TensorType.Builder().indexed("x", 3).build();
+        x.addField(new Field("tensorfield", new TensorDataType(tensorType)));
+        types.registerDocumentType(x);
+    }
+
     private void registerMultiPositionDocumentType() {
         DocumentType x = new DocumentType("testmultipos");
         DataType d = new ArrayDataType(PositionDataType.INSTANCE);
@@ -336,7 +347,6 @@ public class JsonWriterTestCase {
     private Document readDocumentFromJson(String docId, String fields) throws IOException {
         InputStream rawDoc = new ByteArrayInputStream(asFeed(docId, fields));
 
-
         JsonReader r = new JsonReader(types, rawDoc, parserFactory);
         DocumentParseInfo raw = r.parseDocument().get();
         DocumentType docType = r.readDocumentType(raw.documentId);
@@ -433,6 +443,30 @@ public class JsonWriterTestCase {
     }
 
     @Test
+    public void testTensorShortForm() throws IOException {
+        DocumentType documentTypeWithTensor = types.getDocumentType("testindexedtensor");
+        String docId = "id:unittest:testindexedtensor::0";
+        Document doc = new Document(documentTypeWithTensor, docId);
+        Field tensorField = documentTypeWithTensor.getField("tensorfield");
+        Tensor tensor = Tensor.from("tensor(x[3]):[1,2,3]");
+        doc.setFieldValue(tensorField, new TensorFieldValue(tensor));
+
+        assertEqualJson(asDocument(docId, "{ \"tensorfield\": {\"cells\":[{\"address\":{\"x\":\"0\"},\"value\":1.0},{\"address\":{\"x\":\"1\"},\"value\":2.0},{\"address\":{\"x\":\"2\"},\"value\":3.0}]} }"),
+                        writeDocument(doc, false));
+        assertEqualJson(asDocument(docId, "{ \"tensorfield\": {\"type\":\"tensor(x[3])\", \"values\":[1.0, 2.0, 3.0] } }"),
+                        writeDocument(doc, true));
+    }
+
+    private byte[] writeDocument(Document doc, boolean tensorShortForm) throws IOException {
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        JsonFactory factory = new JsonFactory();
+        JsonGenerator generator = factory.createGenerator(out);
+        JsonWriter writer = new JsonWriter(generator, tensorShortForm);
+        writer.write(doc);
+        return out.toByteArray();
+    }
+
+    @Test
     public void non_empty_reference_field_is_roundtrip_json_serialized() throws IOException {
         roundTripEquality("id:unittest:testrefs::helloworld",
                 "{ \"ref_field\": \"id:unittest:smoke::and_mirrors_too\" }");