Support direct tensor rendering

2 files changed, 78 insertions, 47 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/serialization/JsonFormat.java b/vespajlib/src/main/java/com/yahoo/tensor/serialization/JsonFormat.java
index 8322b3b6327..68997c82d3e 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/serialization/JsonFormat.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/serialization/JsonFormat.java
@@ -37,52 +37,79 @@ import java.util.stream.Collectors;
  */
 public class JsonFormat {
 
-    /** Serializes the given tensor value into JSON format */
-    public static byte[] encode(Tensor tensor) {
+    /**
+     * Serializes the given tensor value into JSON format.
+     *
+     * @param tensor the tensor to serialize
+     * @param shortForm whether to encode in a short type-dependent format
+     * @param directValues whether to encode values directly, or wrapped in am object containing "type" and "cells"
+     */
+    public static byte[] encode(Tensor tensor, boolean shortForm, boolean directValues) {
         Slime slime = new Slime();
-        Cursor root = slime.setObject();
-        encodeCells(tensor, root);
+        if (shortForm) {
+            Cursor root = null;
+            if ( ! directValues) {
+                root = slime.setObject();
+                root.setString("type", tensor.type().toString());
+            }
+
+            if (tensor instanceof IndexedTensor denseTensor) {
+                // Encode as nested lists if indexed tensor
+                Cursor parent = root == null ? slime.setArray() : root.setArray("values");
+                encodeValues(denseTensor, parent, new long[denseTensor.dimensionSizes().dimensions()], 0);
+            } else if (tensor instanceof MappedTensor && tensor.type().dimensions().size() == 1) {
+                // Short form for a single mapped dimension
+                Cursor parent = root == null ? slime.setObject() : root.setObject("cells");
+                encodeSingleDimensionCells((MappedTensor) tensor, parent);
+            } else if (tensor instanceof MixedTensor &&
+                       tensor.type().dimensions().stream().anyMatch(TensorType.Dimension::isMapped)) {
+                // Short form for a mixed tensor
+                boolean singleMapped = tensor.type().dimensions().stream().filter(TensorType.Dimension::isMapped).count() == 1;
+                Cursor parent = root == null ? ( singleMapped ? slime.setObject() : slime.setArray() )
+                                             : ( singleMapped ? root.setObject("blocks") : root.setArray("blocks"));
+                encodeBlocks((MixedTensor) tensor, parent);
+            } else {
+                // default to standard cell address output
+                Cursor parent = root == null ? slime.setArray() : root.setArray("cells");
+                encodeCells(tensor, parent);
+            }
+
+            return com.yahoo.slime.JsonFormat.toJsonBytes(slime);
+        }
+        else {
+            Cursor root = slime.setObject();
+            root.setString("type", tensor.type().toString());
+            encodeCells(tensor, root.setArray("cells"));
+        }
         return com.yahoo.slime.JsonFormat.toJsonBytes(slime);
     }
 
-    /** Serializes the given tensor type and value into JSON format */
+    /** Serializes the given tensor value into JSON format, in long format, wrapped in an object containing "cells" only. */
+    public static byte[] encode(Tensor tensor) {
+        return encode(tensor, false, false);
+    }
+
+    /**
+     * Serializes the given tensor type and value into JSON format.
+     *
+     * @deprecated use #encode(#Tensor, boolean, boolean)
+     */
+    @Deprecated // TODO: Remove on Vespa 9
     public static byte[] encodeWithType(Tensor tensor) {
-        Slime slime = new Slime();
-        Cursor root = slime.setObject();
-        root.setString("type", tensor.type().toString());
-        encodeCells(tensor, root);
-        return com.yahoo.slime.JsonFormat.toJsonBytes(slime);
+        return encode(tensor, false, false);
     }
 
-    /** Serializes the given tensor type and value into a short-form JSON format */
+    /**
+     * Serializes the given tensor type and value into a short-form JSON format.
+     *
+     * @deprecated use #encode(#Tensor, boolean, boolean)
+     */
+    @Deprecated // TODO: Remove on Vespa 9
     public static byte[] encodeShortForm(Tensor tensor) {
-        Slime slime = new Slime();
-        Cursor root = slime.setObject();
-        root.setString("type", tensor.type().toString());
-
-        if (tensor instanceof IndexedTensor denseTensor) {
-            // Encode as nested lists if indexed tensor
-            encodeValues(denseTensor, root.setArray("values"), new long[denseTensor.dimensionSizes().dimensions()], 0);
-        }
-        else if (tensor instanceof MappedTensor && tensor.type().dimensions().size() == 1) {
-            // Short form for a single mapped dimension
-            encodeSingleDimensionCells((MappedTensor) tensor, root);
-        }
-        else if (tensor instanceof MixedTensor &&
-                tensor.type().dimensions().stream().filter(TensorType.Dimension::isMapped).count() >= 1) {
-            // Short form for a mixed tensor
-            encodeBlocks((MixedTensor) tensor, root);
-        }
-        else {
-            // default to standard cell address output
-            encodeCells(tensor, root);
-        }
-
-        return com.yahoo.slime.JsonFormat.toJsonBytes(slime);
+        return encode(tensor, true, false);
     }
 
-    private static void encodeCells(Tensor tensor, Cursor rootObject) {
-        Cursor cellsArray = rootObject.setArray("cells");
+    private static void encodeCells(Tensor tensor, Cursor cellsArray) {
         for (Iterator<Tensor.Cell> i = tensor.cellIterator(); i.hasNext(); ) {
             Tensor.Cell cell = i.next();
             Cursor cellObject = cellsArray.addObject();
@@ -91,8 +118,7 @@ public class JsonFormat {
         }
     }
 
-    private static void encodeSingleDimensionCells(MappedTensor tensor, Cursor cursor) {
-        Cursor cells = cursor.setObject("cells");
+    private static void encodeSingleDimensionCells(MappedTensor tensor, Cursor cells) {
         if (tensor.type().dimensions().size() > 1)
             throw new IllegalStateException("JSON encode of mapped tensor can only contain a single dimension");
         tensor.cells().forEach((k,v) -> cells.setDouble(k.label(0), v));
@@ -124,7 +150,6 @@ public class JsonFormat {
         if (mappedDimensions.size() < 1) {
             throw new IllegalArgumentException("Should be ensured by caller");
         }
-        cursor = (mappedDimensions.size() == 1) ? cursor.setObject("blocks") : cursor.setArray("blocks");
 
         // Create tensor type for mapped dimensions subtype
         TensorType mappedSubType = new TensorType.Builder(mappedDimensions).build();
diff --git a/vespajlib/src/test/java/com/yahoo/tensor/serialization/JsonFormatTestCase.java b/vespajlib/src/test/java/com/yahoo/tensor/serialization/JsonFormatTestCase.java
index 1cbf1709be1..4692cf87d59 100644
--- a/vespajlib/src/test/java/com/yahoo/tensor/serialization/JsonFormatTestCase.java
+++ b/vespajlib/src/test/java/com/yahoo/tensor/serialization/JsonFormatTestCase.java
@@ -47,8 +47,8 @@ public class JsonFormatTestCase {
         builder.cell().label("x", "a").label("y", "b").value(2.0);
         builder.cell().label("x", "c").label("y", "d").value(3.0);
         Tensor tensor = builder.build();
-        byte[] json = JsonFormat.encode(tensor);
-        assertEquals("{\"cells\":[" +
+        byte[] json = JsonFormat.encode(tensor, false, false);
+        assertEquals("{\"type\":\"tensor(x{},y{})\",\"cells\":[" +
                      "{\"address\":{\"x\":\"a\",\"y\":\"b\"},\"value\":2.0}," +
                      "{\"address\":{\"x\":\"c\",\"y\":\"d\"},\"value\":3.0}" +
                      "]}",
@@ -61,8 +61,8 @@ public class JsonFormatTestCase {
     public void testEmptySparseTensor() {
         Tensor.Builder builder = Tensor.Builder.of(TensorType.fromSpec("tensor(x{},y{})"));
         Tensor tensor = builder.build();
-        byte[] json = JsonFormat.encode(tensor);
-        assertEquals("{\"cells\":[]}",
+        byte[] json = JsonFormat.encode(tensor, false, false);
+        assertEquals("{\"type\":\"tensor(x{},y{})\",\"cells\":[]}",
                      new String(json, StandardCharsets.UTF_8));
         Tensor decoded = JsonFormat.decode(tensor.type(), json);
         assertEquals(tensor, decoded);
@@ -95,8 +95,8 @@ public class JsonFormatTestCase {
         builder.cell().label("x", 1).label("y", 0).value(5.0);
         builder.cell().label("x", 1).label("y", 1).value(7.0);
         Tensor tensor = builder.build();
-        byte[] json = JsonFormat.encode(tensor);
-        assertEquals("{\"cells\":[" +
+        byte[] json = JsonFormat.encode(tensor, false, false);
+        assertEquals("{\"type\":\"tensor(x[2],y[2])\",\"cells\":[" +
                      "{\"address\":{\"x\":\"0\",\"y\":\"0\"},\"value\":2.0}," +
                      "{\"address\":{\"x\":\"0\",\"y\":\"1\"},\"value\":3.0}," +
                      "{\"address\":{\"x\":\"1\",\"y\":\"0\"},\"value\":5.0}," +
@@ -206,9 +206,14 @@ public class JsonFormatTestCase {
         builder.cell().label("x", 1).label("y", 1).value(0.0);
         builder.cell().label("x", 1).label("y", 2).value(42.0);
         Tensor expected = builder.build();
+
         String denseJson = "{\"values\":\"027FFF80002A\"}";
         Tensor decoded = JsonFormat.decode(expected.type(), denseJson.getBytes(StandardCharsets.UTF_8));
         assertEquals(expected, decoded);
+
+        denseJson = "\"027FFF80002A\"";
+        decoded = JsonFormat.decode(expected.type(), denseJson.getBytes(StandardCharsets.UTF_8));
+        assertEquals(expected, decoded);
     }
 
     @Test
@@ -233,6 +238,7 @@ public class JsonFormatTestCase {
         builder.cell().label("x", 1).label("y", 1).value(6.0);
         builder.cell().label("x", 1).label("y", 2).value(7.0);
         Tensor expected = builder.build();
+
         String mixedJson = "{\"blocks\":[" +
                            "{\"address\":{\"x\":\"0\"},\"values\":\"020304\"}," +
                            "{\"address\":{\"x\":\"1\"},\"values\":\"050607\"}" +
@@ -375,7 +381,7 @@ public class JsonFormatTestCase {
     }
 
     private void assertEncodeDecode(Tensor tensor) {
-        Tensor decoded = JsonFormat.decode(tensor.type(), JsonFormat.encodeWithType(tensor));
+        Tensor decoded = JsonFormat.decode(tensor.type(), JsonFormat.encode(tensor, false, false));
         assertEquals(tensor, decoded);
         assertEquals(tensor.type(), decoded.type());
     }
@@ -403,7 +409,7 @@ public class JsonFormatTestCase {
     }
 
     private void assertEncodeShortForm(Tensor tensor, String expected) {
-        byte[] json = JsonFormat.encodeShortForm(tensor);
+        byte[] json = JsonFormat.encode(tensor, true, false);
         assertEquals(expected, new String(json, StandardCharsets.UTF_8));
     }