1 files changed, 58 insertions, 27 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 bebd706f815..87157495485 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/serialization/JsonFormat.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/serialization/JsonFormat.java
@@ -16,9 +16,11 @@ import com.yahoo.tensor.MixedTensor;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorAddress;
 import com.yahoo.tensor.TensorType;
+import com.yahoo.tensor.evaluation.Name;
 import com.yahoo.tensor.functions.ConstantTensor;
 import com.yahoo.tensor.functions.Slice;
 
+import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
@@ -61,23 +63,23 @@ public class JsonFormat {
         // Encode as nested lists if indexed tensor
         if (tensor instanceof IndexedTensor) {
             IndexedTensor denseTensor = (IndexedTensor) tensor;
-            encodeList(denseTensor, root.setArray("value"), new long[denseTensor.dimensionSizes().dimensions()], 0);
+            encodeValues(denseTensor, root.setArray("values"), new long[denseTensor.dimensionSizes().dimensions()], 0);
         }
 
         // Short form for a single mapped dimension
         else if (tensor instanceof MappedTensor && tensor.type().dimensions().size() == 1) {
-            encodeMap((MappedTensor) tensor, root.setObject("value"));
+            encodeSingleDimensionCells((MappedTensor) tensor, root);
         }
 
-        // Short form for a mixed tensor with a single mapped dimension
+        // Short form for a mixed tensor
         else if (tensor instanceof MixedTensor &&
-                tensor.type().dimensions().stream().filter(TensorType.Dimension::isMapped).count() == 1) {
-            encodeMapBlocks((MixedTensor) tensor, root.setObject("value"));
+                tensor.type().dimensions().stream().filter(TensorType.Dimension::isMapped).count() >= 1) {
+            encodeBlocks((MixedTensor) tensor, root);
         }
 
         // No other short forms exist: default to standard cell address output
         else {
-            encodeCells(tensor, root.setObject("value"));
+            encodeCells(tensor, root);
         }
 
         return com.yahoo.slime.JsonFormat.toJsonBytes(slime);
@@ -93,47 +95,76 @@ public class JsonFormat {
         }
     }
 
+    private static void encodeSingleDimensionCells(MappedTensor tensor, Cursor cursor) {
+        Cursor cells = cursor.setObject("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));
+    }
+
     private static void encodeAddress(TensorType type, TensorAddress address, Cursor addressObject) {
         for (int i = 0; i < address.size(); i++)
             addressObject.setString(type.dimensions().get(i).name(), address.label(i));
     }
 
-    private static void encodeList(IndexedTensor tensor, Cursor cursor, long[] indexes, int dimension) {
+    private static void encodeValues(IndexedTensor tensor, Cursor cursor, long[] indexes, int dimension) {
         DimensionSizes sizes = tensor.dimensionSizes();
         for (indexes[dimension] = 0; indexes[dimension] < sizes.size(dimension); ++indexes[dimension]) {
             if (dimension < (sizes.dimensions() - 1)) {
-                encodeList(tensor, cursor.addArray(), indexes, dimension + 1);
+                encodeValues(tensor, cursor.addArray(), indexes, dimension + 1);
             } else {
                 cursor.addDouble(tensor.get(indexes));
             }
         }
     }
 
-    private static void encodeMap(MappedTensor tensor, Cursor cursor) {
-        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) -> cursor.setDouble(k.label(0), v));
-    }
-
-    private static void encodeMapBlocks(MixedTensor tensor, Cursor cursor) {
-        var mappedDimensions = tensor.type().dimensions().stream().filter(d -> !d.isIndexed()).collect(Collectors.toList());
-        if (mappedDimensions.size() != 1) {
+    private static void encodeBlocks(MixedTensor tensor, Cursor cursor) {
+        var mappedDimensions = tensor.type().dimensions().stream().filter(d -> d.isMapped())
+                .map(d -> TensorType.Dimension.mapped(d.name())).collect(Collectors.toList());
+        if (mappedDimensions.size() < 1) {
             throw new IllegalArgumentException("Should be ensured by caller");
         }
-        String mappedDimensionName = mappedDimensions.get(0).name();
-        int mappedDimensionIndex = tensor.type().indexOfDimension(mappedDimensionName).
-                orElseThrow(() -> new IllegalStateException("Could not find mapped dimension index"));
+        cursor = (mappedDimensions.size() == 1) ? cursor.setObject("blocks") : cursor.setArray("blocks");
 
-        // Find all unique indices for the mapped dimension
-        Set<String> mappedIndices = new HashSet<>();
-        tensor.cellIterator().forEachRemaining((cell) -> mappedIndices.add(cell.getKey().label(mappedDimensionIndex)));
+        // Create tensor type for mapped dimensions subtype
+        TensorType mappedSubType = new TensorType.Builder(mappedDimensions).build();
+
+        // Find all unique indices for the mapped dimensions
+        Set<TensorAddress> denseSubSpaceAddresses = new HashSet<>();
+        tensor.cellIterator().forEachRemaining((cell) -> {
+            denseSubSpaceAddresses.add(subAddress(cell.getKey(), mappedSubType, tensor.type()));
+        });
 
         // Slice out dense subspace of each and encode dense subspace as a list
-        for (String mappedIndex : mappedIndices) {
-            IndexedTensor denseSubspace = (IndexedTensor) new Slice<>(new ConstantTensor<>(tensor),
-                    List.of(new Slice.DimensionValue<>(mappedDimensionName, mappedIndex))).evaluate();
-            encodeList(denseSubspace, cursor.setArray(mappedIndex), new long[denseSubspace.dimensionSizes().dimensions()], 0);
+        for (TensorAddress denseSubSpaceAddress : denseSubSpaceAddresses) {
+            IndexedTensor denseSubspace = (IndexedTensor) sliceSubAddress(tensor, denseSubSpaceAddress, mappedSubType);
+
+            if (mappedDimensions.size() == 1) {
+                encodeValues(denseSubspace, cursor.setArray(denseSubSpaceAddress.label(0)), new long[denseSubspace.dimensionSizes().dimensions()], 0);
+            } else {
+                Cursor block = cursor.addObject();
+                encodeAddress(mappedSubType, denseSubSpaceAddress, block.setObject("address"));
+                encodeValues(denseSubspace, block.setArray("values"), new long[denseSubspace.dimensionSizes().dimensions()], 0);
+            }
+
+        }
+    }
+
+    private static TensorAddress subAddress(TensorAddress address, TensorType subType, TensorType origType) {
+        TensorAddress.Builder builder = new TensorAddress.Builder(subType);
+        for (TensorType.Dimension dim : subType.dimensions()) {
+            builder.add(dim.name(), address.label(origType.indexOfDimension(dim.name()).
+                    orElseThrow(() -> new IllegalStateException("Could not find mapped dimension index"))));
+        }
+        return builder.build();
+    }
+
+    private static Tensor sliceSubAddress(Tensor tensor, TensorAddress subAddress, TensorType subType) {
+        List<Slice.DimensionValue<Name>> sliceDims = new ArrayList<>(subAddress.size());
+        for (int i = 0; i < subAddress.size(); ++i) {
+            sliceDims.add(new Slice.DimensionValue<>(subType.dimensions().get(i).name(), subAddress.label(i)));
         }
+        return new Slice<>(new ConstantTensor<>(tensor), sliceDims).evaluate();
     }
 
     /** Deserializes the given tensor from JSON format */