Initial skeleton for type information in serialization format.

4 files changed, 80 insertions, 3 deletions

diff --git a/vespajlib/abi-spec.json b/vespajlib/abi-spec.json
index b450b183c08..766b0daedcc 100644
--- a/vespajlib/abi-spec.json
+++ b/vespajlib/abi-spec.json
@@ -1267,6 +1267,23 @@
     ],
     "fields": []
   },
+  "com.yahoo.tensor.TensorType$ValueType": {
+    "superClass": "java.lang.Enum",
+    "interfaces": [],
+    "attributes": [
+      "public",
+      "final",
+      "enum"
+    ],
+    "methods": [
+      "public static com.yahoo.tensor.TensorType$ValueType[] values()",
+      "public static com.yahoo.tensor.TensorType$ValueType valueOf(java.lang.String)"
+    ],
+    "fields": [
+      "public static final enum com.yahoo.tensor.TensorType$ValueType DOUBLE",
+      "public static final enum com.yahoo.tensor.TensorType$ValueType FLOAT"
+    ]
+  },
   "com.yahoo.tensor.TensorType": {
     "superClass": "java.lang.Object",
     "interfaces": [],
@@ -1274,6 +1291,7 @@
       "public"
     ],
     "methods": [
+      "public final com.yahoo.tensor.TensorType$ValueType valueType()",
       "public static com.yahoo.tensor.TensorType fromSpec(java.lang.String)",
       "public int rank()",
       "public java.util.List dimensions()",
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java b/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
index acba9eafd71..bb7a976af9b 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java
@@ -24,9 +24,15 @@ import java.util.stream.Collectors;
  */
 public class TensorType {
 
+    public enum ValueType { DOUBLE, FLOAT};
+
     /** The empty tensor type - which is the same as a double */
     public static final TensorType empty = new TensorType(Collections.emptyList());
 
+    private final ValueType valueType = ValueType.DOUBLE;
+
+    public final ValueType valueType() { return valueType; }
+
     /** Sorted list of the dimensions of this */
     private final ImmutableList<Dimension> dimensions;
 
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/serialization/DenseBinaryFormat.java b/vespajlib/src/main/java/com/yahoo/tensor/serialization/DenseBinaryFormat.java
index 46d533cc6bd..6dcc870ef94 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/serialization/DenseBinaryFormat.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/serialization/DenseBinaryFormat.java
@@ -22,14 +22,40 @@ import java.util.Optional;
  */
 public class DenseBinaryFormat implements BinaryFormat {
 
+    static private final int DOUBLE_VALUE_TYPE = 0; // Not encoded as it is default, and you know the type when deserializing
+    static private final int FLOAT_VALUE_TYPE = 1;
+
+    enum EncodeType {NO_DEFAULT, DOUBLE_IS_DEFAULT}
+    private final EncodeType encodeType;
+    DenseBinaryFormat() {
+        encodeType = EncodeType.DOUBLE_IS_DEFAULT;
+    }
+    DenseBinaryFormat(EncodeType encodeType) {
+        this.encodeType = encodeType;
+    }
+
     @Override
     public void encode(GrowableByteBuffer buffer, Tensor tensor) {
         if ( ! ( tensor instanceof IndexedTensor))
             throw new RuntimeException("The dense format is only supported for indexed tensors");
+        encodeValueType(buffer, tensor.type().valueType());
         encodeDimensions(buffer, (IndexedTensor)tensor);
         encodeCells(buffer, tensor);
     }
 
+    private void encodeValueType(GrowableByteBuffer buffer, TensorType.ValueType valueType) {
+        switch (valueType) {
+            case DOUBLE:
+                if (encodeType != EncodeType.DOUBLE_IS_DEFAULT) {
+                    buffer.putInt1_4Bytes(DOUBLE_VALUE_TYPE);
+                }
+                break;
+            case FLOAT:
+                buffer.putInt1_4Bytes(FLOAT_VALUE_TYPE);
+                break;
+        }
+    }
+
     private void encodeDimensions(GrowableByteBuffer buffer, IndexedTensor tensor) {
         buffer.putInt1_4Bytes(tensor.type().dimensions().size());
         for (int i = 0; i < tensor.type().dimensions().size(); i++) {
@@ -39,11 +65,28 @@ public class DenseBinaryFormat implements BinaryFormat {
     }
 
     private void encodeCells(GrowableByteBuffer buffer, Tensor tensor) {
+        switch (tensor.type().valueType()) {
+            case DOUBLE:
+                encodeCellsAsDouble(buffer, tensor);
+                break;
+            case FLOAT:
+                encodeCellsAsFloat(buffer, tensor);
+                break;
+        }
+    }
+
+    private void encodeCellsAsDouble(GrowableByteBuffer buffer, Tensor tensor) {
         Iterator<Double> i = tensor.valueIterator();
         while (i.hasNext())
             buffer.putDouble(i.next());
     }
 
+    private void encodeCellsAsFloat(GrowableByteBuffer buffer, Tensor tensor) {
+        Iterator<Double> i = tensor.valueIterator();
+        while (i.hasNext())
+            buffer.putFloat(i.next().floatValue());
+    }
+
     @Override
     public Tensor decode(Optional<TensorType> optionalType, GrowableByteBuffer buffer) {
         TensorType type;
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/serialization/TypedBinaryFormat.java b/vespajlib/src/main/java/com/yahoo/tensor/serialization/TypedBinaryFormat.java
index 1c18adc8652..f2c5d4e2bd8 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/serialization/TypedBinaryFormat.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/serialization/TypedBinaryFormat.java
@@ -23,6 +23,7 @@ public class TypedBinaryFormat {
     private static final int SPARSE_BINARY_FORMAT_TYPE = 1;
     private static final int DENSE_BINARY_FORMAT_TYPE = 2;
     private static final int MIXED_BINARY_FORMAT_TYPE = 3;
+    private static final int TYPED_DENSE_BINARY_FORMAT_TYPE = 4;
 
     public static byte[] encode(Tensor tensor) {
         GrowableByteBuffer buffer = new GrowableByteBuffer();
@@ -31,8 +32,16 @@ public class TypedBinaryFormat {
             new MixedBinaryFormat().encode(buffer, tensor);
         }
         else if (tensor instanceof IndexedTensor) {
-            buffer.putInt1_4Bytes(DENSE_BINARY_FORMAT_TYPE);
-            new DenseBinaryFormat().encode(buffer, tensor);
+            switch (tensor.type().valueType()) {
+                case DOUBLE:
+                    buffer.putInt1_4Bytes(DENSE_BINARY_FORMAT_TYPE);
+                    new DenseBinaryFormat(DenseBinaryFormat.EncodeType.DOUBLE_IS_DEFAULT).encode(buffer, tensor);
+                    break;
+                default:
+                    buffer.putInt1_4Bytes(TYPED_DENSE_BINARY_FORMAT_TYPE);
+                    new DenseBinaryFormat(DenseBinaryFormat.EncodeType.NO_DEFAULT).encode(buffer, tensor);
+                    break;
+            }
         }
         else {
             buffer.putInt1_4Bytes(SPARSE_BINARY_FORMAT_TYPE);
@@ -57,7 +66,8 @@ public class TypedBinaryFormat {
         switch (formatType) {
             case MIXED_BINARY_FORMAT_TYPE: return new MixedBinaryFormat().decode(type, buffer);
             case SPARSE_BINARY_FORMAT_TYPE: return new SparseBinaryFormat().decode(type, buffer);
-            case DENSE_BINARY_FORMAT_TYPE: return new DenseBinaryFormat().decode(type, buffer);
+            case DENSE_BINARY_FORMAT_TYPE: return new DenseBinaryFormat(DenseBinaryFormat.EncodeType.DOUBLE_IS_DEFAULT).decode(type, buffer);
+            case TYPED_DENSE_BINARY_FORMAT_TYPE: return new DenseBinaryFormat(DenseBinaryFormat.EncodeType.NO_DEFAULT).decode(type, buffer);
             default: throw new IllegalArgumentException("Binary format type " + formatType + " is unknown");
         }
     }