allow short-form JSON for 1-d constants

4 files changed, 101 insertions, 12 deletions

diff --git a/eval/src/tests/eval/value_cache/dense-short.json b/eval/src/tests/eval/value_cache/dense-short.json
new file mode 100644
index 00000000000..40121135544
--- /dev/null
+++ b/eval/src/tests/eval/value_cache/dense-short.json
@@ -0,0 +1,3 @@
+{
+  "values": [ 1, 2.0, 3.5 ]
+}
diff --git a/eval/src/tests/eval/value_cache/sparse-short.json b/eval/src/tests/eval/value_cache/sparse-short.json
new file mode 100644
index 00000000000..f10b1b6f9fb
--- /dev/null
+++ b/eval/src/tests/eval/value_cache/sparse-short.json
@@ -0,0 +1,7 @@
+{
+  "cells": {
+    "foo": 1.0,
+    "bar": 2.0,
+    "three": 3.0
+  }
+}
diff --git a/eval/src/tests/eval/value_cache/tensor_loader_test.cpp b/eval/src/tests/eval/value_cache/tensor_loader_test.cpp
index 1a77cfe847b..ed35b8b108f 100644
--- a/eval/src/tests/eval/value_cache/tensor_loader_test.cpp
+++ b/eval/src/tests/eval/value_cache/tensor_loader_test.cpp
@@ -19,12 +19,26 @@ TensorSpec make_dense_tensor() {
         .add({{"x", 1}, {"y", 1}}, 4.0);
 }
 
+TensorSpec make_simple_dense_tensor() {
+    return TensorSpec("tensor(z[3])")
+        .add({{"z", 0}}, 1.0)
+        .add({{"z", 1}}, 2.0)
+        .add({{"z", 2}}, 3.5);
+}
+
 TensorSpec make_sparse_tensor() {
     return TensorSpec("tensor(x{},y{})")
         .add({{"x", "foo"}, {"y", "bar"}}, 1.0)
         .add({{"x", "bar"}, {"y", "foo"}}, 2.0);
 }
 
+TensorSpec make_simple_sparse_tensor() {
+    return TensorSpec("tensor(mydim{})")
+        .add({{"mydim", "foo"}}, 1.0)
+        .add({{"mydim", "three"}}, 3.0)
+        .add({{"mydim", "bar"}}, 2.0);
+}
+
 TensorSpec make_mixed_tensor() {
     return TensorSpec("tensor(x{},y[2])")
         .add({{"x", "foo"}, {"y", 0}}, 1.0)
@@ -75,6 +89,14 @@ TEST_F("require that lz4 compressed sparse tensor can be loaded", ConstantTensor
     TEST_DO(verify_tensor(make_sparse_tensor(), f1.create(TEST_PATH("sparse.json.lz4"), "tensor(x{},y{})")));
 }
 
+TEST_F("require that sparse tensor short form can be loaded", ConstantTensorLoader(factory)) {
+    TEST_DO(verify_tensor(make_simple_sparse_tensor(), f1.create(TEST_PATH("sparse-short.json"), "tensor(mydim{})")));
+}
+
+TEST_F("require that dense tensor short form can be loaded", ConstantTensorLoader(factory)) {
+    TEST_DO(verify_tensor(make_simple_dense_tensor(), f1.create(TEST_PATH("dense-short.json"), "tensor(z[3])")));
+}
+
 TEST_F("require that bad lz4 file fails to load creating empty result", ConstantTensorLoader(factory)) {
     TEST_DO(verify_tensor(sparse_tensor_nocells(), f1.create(TEST_PATH("bad_lz4.json.lz4"), "tensor(x{},y{})")));
 }
diff --git a/eval/src/vespa/eval/eval/value_cache/constant_tensor_loader.cpp b/eval/src/vespa/eval/eval/value_cache/constant_tensor_loader.cpp
index 9af473f1f94..a99e6f3b0dd 100644
--- a/eval/src/vespa/eval/eval/value_cache/constant_tensor_loader.cpp
+++ b/eval/src/vespa/eval/eval/value_cache/constant_tensor_loader.cpp
@@ -41,6 +41,52 @@ struct AddressExtractor : ObjectTraverser {
     }
 };
 
+struct SingleMappedExtractor : ObjectTraverser {
+    const vespalib::string &dimension;
+    TensorSpec &spec;
+    SingleMappedExtractor(const vespalib::string &dimension_in, TensorSpec &spec_in)
+        : dimension(dimension_in),
+          spec(spec_in)
+    {}
+    void field(const Memory &symbol, const Inspector &inspector) override {
+        vespalib::string label = symbol.make_string();
+        double value = inspector.asDouble();
+        TensorSpec::Address address;
+        address.emplace(dimension, label);
+        spec.add(address, value);
+    }
+};
+
+
+void decodeSingleMappedForm(const Inspector &root, const ValueType &value_type, TensorSpec &spec) {
+    auto extractor = SingleMappedExtractor(value_type.dimensions()[0].name, spec);
+    root.traverse(extractor);
+}
+
+void decodeSingleDenseForm(const Inspector &values, const ValueType &value_type, TensorSpec &spec) {
+    const auto &dimension = value_type.dimensions()[0].name;
+    for (size_t i = 0; i < values.entries(); ++i) {
+        TensorSpec::Address address;
+        address.emplace(dimension, TensorSpec::Label(i));
+        spec.add(address, values[i].asDouble());
+    }
+}
+
+void decodeLiteralForm(const Inspector &cells, const ValueType &value_type, TensorSpec &spec) {
+    std::set<vespalib::string> indexed;
+    for (const auto &dimension: value_type.dimensions()) {
+        if (dimension.is_indexed()) {
+            indexed.insert(dimension.name);
+        }
+    }
+    for (size_t i = 0; i < cells.entries(); ++i) {
+        TensorSpec::Address address;
+        AddressExtractor extractor(indexed, address);
+        cells[i]["address"].traverse(extractor);
+        spec.add(address, cells[i]["value"].asDouble());
+    }
+}
+
 void decode_json(const vespalib::string &path, Input &input, Slime &slime) {
     if (slime::JsonFormat::decode(input, slime) == 0) {
         LOG(warning, "file contains invalid json: %s", path.c_str());
@@ -90,19 +136,30 @@ ConstantTensorLoader::create(const vespalib::string &path, const vespalib::strin
     }
     Slime slime;
     decode_json(path, slime);
-    std::set<vespalib::string> indexed;
-    for (const auto &dimension: value_type.dimensions()) {
-        if (dimension.is_indexed()) {
-            indexed.insert(dimension.name);
-        }
-    }
     TensorSpec spec(type);
-    const Inspector &cells = slime.get()["cells"];
-    for (size_t i = 0; i < cells.entries(); ++i) {
-        TensorSpec::Address address;
-        AddressExtractor extractor(indexed, address);
-        cells[i]["address"].traverse(extractor);
-        spec.add(address, cells[i]["value"].asDouble());
+    const Inspector &root = slime.get();
+    const Inspector &cells = root["cells"];
+    const Inspector &values = root["values"];
+    if (cells.type().getId() == vespalib::slime::ARRAY::ID) {
+        decodeLiteralForm(cells, value_type, spec);
+    }
+    else if (cells.type().getId() == vespalib::slime::OBJECT::ID
+             && value_type.is_sparse()
+             && value_type.count_mapped_dimensions() == 1)
+    {
+        decodeSingleMappedForm(cells, value_type, spec);
+    }
+    else if (values.type().getId() == vespalib::slime::ARRAY::ID
+             && value_type.is_dense()
+             && value_type.count_indexed_dimensions() == 1)
+    {
+        decodeSingleDenseForm(values, value_type, spec);
+    }
+    else if (root.type().getId() == vespalib::slime::OBJECT::ID
+             && value_type.is_sparse()
+             && value_type.count_mapped_dimensions() == 1)
+    {
+        decodeSingleMappedForm(root, value_type, spec);
     }
     try {
         return std::make_unique<SimpleConstantValue>(value_from_spec(spec, _factory));