diff options
Diffstat (limited to 'vespajlib/src/main/java/com/yahoo/tensor/serialization/JsonFormat.java')
-rw-r--r-- | vespajlib/src/main/java/com/yahoo/tensor/serialization/JsonFormat.java | 23 |
1 files changed, 18 insertions, 5 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 b7e6e67ce73..45e581d73e8 100644 --- a/vespajlib/src/main/java/com/yahoo/tensor/serialization/JsonFormat.java +++ b/vespajlib/src/main/java/com/yahoo/tensor/serialization/JsonFormat.java @@ -25,7 +25,6 @@ import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Set; -import java.util.stream.Collectors; /** * Writes tensors on the JSON format used in Vespa tensor document fields: @@ -113,14 +112,14 @@ public class JsonFormat { Tensor.Cell cell = i.next(); Cursor cellObject = cellsArray.addObject(); encodeAddress(tensor.type(), cell.getKey(), cellObject.setObject("address")); - cellObject.setDouble("value", cell.getValue()); + setValue("value", cell.getValue(), tensor.type().valueType(), cellObject); } } 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)); + tensor.cells().forEach((k,v) -> setValue(k.label(0), v, tensor.type().valueType(), cells)); } private static void encodeAddress(TensorType type, TensorAddress address, Cursor addressObject) { @@ -131,13 +130,13 @@ public class JsonFormat { private static void encodeValues(IndexedTensor tensor, Cursor cursor, long[] indexes, int dimension) { DimensionSizes sizes = tensor.dimensionSizes(); if (indexes.length == 0) { - cursor.addDouble(tensor.get(0)); + addValue(tensor.get(0), tensor.type().valueType(), cursor); } else { for (indexes[dimension] = 0; indexes[dimension] < sizes.size(dimension); ++indexes[dimension]) { if (dimension < (sizes.dimensions() - 1)) { encodeValues(tensor, cursor.addArray(), indexes, dimension + 1); } else { - cursor.addDouble(tensor.get(indexes)); + addValue(tensor.get(indexes), tensor.type().valueType(), cursor); } } } @@ -174,6 +173,20 @@ public class JsonFormat { } } + private static void addValue(double value, TensorType.Value valueType, Cursor cursor) { + if (valueType == TensorType.Value.INT8) + cursor.addLong((long)value); + else + cursor.addDouble(value); + } + + private static void setValue(String field, double value, TensorType.Value valueType, Cursor cursor) { + if (valueType == TensorType.Value.INT8) + cursor.setLong(field, (long)value); + else + cursor.setDouble(field, value); + } + private static TensorAddress subAddress(TensorAddress address, TensorType subType, TensorType origType) { TensorAddress.Builder builder = new TensorAddress.Builder(subType); for (TensorType.Dimension dim : subType.dimensions()) { |