support binary formats with cell type in reference implementation

1 files changed, 208 insertions, 0 deletions

diff --git a/eval/src/tests/eval/simple_tensor/simple_tensor_test.cpp b/eval/src/tests/eval/simple_tensor/simple_tensor_test.cpp
index c3b42124155..aa4c3b8c021 100644
--- a/eval/src/tests/eval/simple_tensor/simple_tensor_test.cpp
+++ b/eval/src/tests/eval/simple_tensor/simple_tensor_test.cpp
@@ -4,6 +4,8 @@
 #include <vespa/eval/eval/simple_tensor_engine.h>
 #include <vespa/eval/eval/operation.h>
 #include <vespa/vespalib/util/stash.h>
+#include <vespa/vespalib/data/memory.h>
+#include <vespa/vespalib/objects/nbostream.h>
 #include <iostream>
 
 using namespace vespalib::eval;
@@ -12,6 +14,8 @@ using Cell = SimpleTensor::Cell;
 using Cells = SimpleTensor::Cells;
 using Address = SimpleTensor::Address;
 using Stash = vespalib::Stash;
+using vespalib::nbostream;
+using vespalib::Memory;
 
 TensorSpec to_spec(const Value &a) { return SimpleTensorEngine::ref().to_spec(a); }
 
@@ -143,4 +147,208 @@ TEST("require that simple tensors support dimension reduction") {
     EXPECT_NOT_EQUAL(to_spec(*result_sum_y), to_spec(*result_sum_x));
 }
 
+//-----------------------------------------------------------------------------
+
+struct SparseTensorExample {
+    TensorSpec make_spec() const {
+        return TensorSpec("tensor(x{},y{})")
+            .add({{"x","a"},{"y","a"}}, 1)
+            .add({{"x","a"},{"y","b"}}, 2)
+            .add({{"x","b"},{"y","a"}}, 3);
+    }
+    std::unique_ptr<SimpleTensor> make_tensor() const {
+        return SimpleTensor::create(make_spec());
+    }
+    template <typename T>
+    void encode_inner(nbostream &dst) const {
+        dst.putInt1_4Bytes(2);
+        dst.writeSmallString("x");
+        dst.writeSmallString("y");
+        dst.putInt1_4Bytes(3);
+        dst.writeSmallString("a");
+        dst.writeSmallString("a");
+        dst << (T) 1;
+        dst.writeSmallString("a");
+        dst.writeSmallString("b");
+        dst << (T) 2;
+        dst.writeSmallString("b");
+        dst.writeSmallString("a");
+        dst << (T) 3;
+    }
+    void encode_default(nbostream &dst) const {
+        dst.putInt1_4Bytes(1);
+        encode_inner<double>(dst);
+    }
+    void encode_with_double(nbostream &dst) const {
+        dst.putInt1_4Bytes(5);
+        dst.putInt1_4Bytes(0);
+        encode_inner<double>(dst);
+    }
+    void encode_with_float(nbostream &dst) const {
+        dst.putInt1_4Bytes(5);
+        dst.putInt1_4Bytes(1);
+        encode_inner<float>(dst);
+    }
+};
+
+TEST_F("require that sparse tensors can be decoded", SparseTensorExample()) {
+    nbostream data1;
+    nbostream data2;
+    nbostream data3;
+    f1.encode_default(data1);
+    f1.encode_with_double(data2);
+    f1.encode_with_float(data3);
+    EXPECT_EQUAL(to_spec(*SimpleTensor::decode(data1)), f1.make_spec());
+    EXPECT_EQUAL(to_spec(*SimpleTensor::decode(data2)), f1.make_spec());
+    EXPECT_EQUAL(to_spec(*SimpleTensor::decode(data3)), f1.make_spec());
+}
+
+TEST_F("require that sparse tensors can be encoded", SparseTensorExample()) {
+    nbostream data;
+    nbostream expect;
+    SimpleTensor::encode(*f1.make_tensor(), data);
+    f1.encode_default(expect);
+    EXPECT_EQUAL(Memory(data.peek(), data.size()), Memory(expect.peek(), expect.size()));
+}
+
+//-----------------------------------------------------------------------------
+
+struct DenseTensorExample {
+    TensorSpec make_spec() const {
+        return TensorSpec("tensor(x[3],y[2])")
+            .add({{"x",0},{"y",0}}, 1)
+            .add({{"x",0},{"y",1}}, 2)
+            .add({{"x",1},{"y",0}}, 3)
+            .add({{"x",1},{"y",1}}, 4)
+            .add({{"x",2},{"y",0}}, 5)
+            .add({{"x",2},{"y",1}}, 6);
+    }
+    std::unique_ptr<SimpleTensor> make_tensor() const {
+        return SimpleTensor::create(make_spec());
+    }
+    template <typename T>
+    void encode_inner(nbostream &dst) const {
+        dst.putInt1_4Bytes(2);
+        dst.writeSmallString("x");
+        dst.putInt1_4Bytes(3);
+        dst.writeSmallString("y");
+        dst.putInt1_4Bytes(2);
+        dst << (T) 1;
+        dst << (T) 2;
+        dst << (T) 3;
+        dst << (T) 4;
+        dst << (T) 5;
+        dst << (T) 6;
+    }
+    void encode_default(nbostream &dst) const {
+        dst.putInt1_4Bytes(2);
+        encode_inner<double>(dst);
+    }
+    void encode_with_double(nbostream &dst) const {
+        dst.putInt1_4Bytes(6);
+        dst.putInt1_4Bytes(0);
+        encode_inner<double>(dst);
+    }
+    void encode_with_float(nbostream &dst) const {
+        dst.putInt1_4Bytes(6);
+        dst.putInt1_4Bytes(1);
+        encode_inner<float>(dst);
+    }
+};
+
+TEST_F("require that dense tensors can be decoded", DenseTensorExample()) {
+    nbostream data1;
+    nbostream data2;
+    nbostream data3;
+    f1.encode_default(data1);
+    f1.encode_with_double(data2);
+    f1.encode_with_float(data3);
+    EXPECT_EQUAL(to_spec(*SimpleTensor::decode(data1)), f1.make_spec());
+    EXPECT_EQUAL(to_spec(*SimpleTensor::decode(data2)), f1.make_spec());
+    EXPECT_EQUAL(to_spec(*SimpleTensor::decode(data3)), f1.make_spec());
+}
+
+TEST_F("require that dense tensors can be encoded", DenseTensorExample()) {
+    nbostream data;
+    nbostream expect;
+    SimpleTensor::encode(*f1.make_tensor(), data);
+    f1.encode_default(expect);
+    EXPECT_EQUAL(Memory(data.peek(), data.size()), Memory(expect.peek(), expect.size()));
+}
+
+//-----------------------------------------------------------------------------
+
+struct MixedTensorExample {
+    TensorSpec make_spec() const {
+        return TensorSpec("tensor(x{},y{},z[2])")
+            .add({{"x","a"},{"y","a"},{"z",0}}, 1)
+            .add({{"x","a"},{"y","a"},{"z",1}}, 2)
+            .add({{"x","a"},{"y","b"},{"z",0}}, 3)
+            .add({{"x","a"},{"y","b"},{"z",1}}, 4)
+            .add({{"x","b"},{"y","a"},{"z",0}}, 5)
+            .add({{"x","b"},{"y","a"},{"z",1}}, 6);
+    }
+    std::unique_ptr<SimpleTensor> make_tensor() const {
+        return SimpleTensor::create(make_spec());
+    }
+    template <typename T>
+    void encode_inner(nbostream &dst) const {
+        dst.putInt1_4Bytes(2);
+        dst.writeSmallString("x");
+        dst.writeSmallString("y");
+        dst.putInt1_4Bytes(1);
+        dst.writeSmallString("z");
+        dst.putInt1_4Bytes(2);
+        dst.putInt1_4Bytes(3);
+        dst.writeSmallString("a");
+        dst.writeSmallString("a");
+        dst << (T) 1;
+        dst << (T) 2;
+        dst.writeSmallString("a");
+        dst.writeSmallString("b");
+        dst << (T) 3;
+        dst << (T) 4;
+        dst.writeSmallString("b");
+        dst.writeSmallString("a");
+        dst << (T) 5;
+        dst << (T) 6;
+    }
+    void encode_default(nbostream &dst) const {
+        dst.putInt1_4Bytes(3);
+        encode_inner<double>(dst);
+    }
+    void encode_with_double(nbostream &dst) const {
+        dst.putInt1_4Bytes(7);
+        dst.putInt1_4Bytes(0);
+        encode_inner<double>(dst);
+    }
+    void encode_with_float(nbostream &dst) const {
+        dst.putInt1_4Bytes(7);
+        dst.putInt1_4Bytes(1);
+        encode_inner<float>(dst);
+    }
+};
+
+TEST_F("require that mixed tensors can be decoded", MixedTensorExample()) {
+    nbostream data1;
+    nbostream data2;
+    nbostream data3;
+    f1.encode_default(data1);
+    f1.encode_with_double(data2);
+    f1.encode_with_float(data3);
+    EXPECT_EQUAL(to_spec(*SimpleTensor::decode(data1)), f1.make_spec());
+    EXPECT_EQUAL(to_spec(*SimpleTensor::decode(data2)), f1.make_spec());
+    EXPECT_EQUAL(to_spec(*SimpleTensor::decode(data3)), f1.make_spec());
+}
+
+TEST_F("require that mixed tensors can be encoded", MixedTensorExample()) {
+    nbostream data;
+    nbostream expect;
+    SimpleTensor::encode(*f1.make_tensor(), data);
+    f1.encode_default(expect);
+    EXPECT_EQUAL(Memory(data.peek(), data.size()), Memory(expect.peek(), expect.size()));
+}
+
+//-----------------------------------------------------------------------------
+
 TEST_MAIN() { TEST_RUN_ALL(); }