add support for loading lz4 compressed json tensors

5 files changed, 39 insertions, 10 deletions

diff --git a/eval/src/tests/eval/value_cache/bad_lz4.json.lz4 b/eval/src/tests/eval/value_cache/bad_lz4.json.lz4
new file mode 100644
index 00000000000..e17d91251a3
--- /dev/null
+++ b/eval/src/tests/eval/value_cache/bad_lz4.json.lz4
diff --git a/eval/src/tests/eval/value_cache/dense.json.lz4 b/eval/src/tests/eval/value_cache/dense.json.lz4
new file mode 100644
index 00000000000..166d74ed6f6
--- /dev/null
+++ b/eval/src/tests/eval/value_cache/dense.json.lz4
diff --git a/eval/src/tests/eval/value_cache/sparse.json.lz4 b/eval/src/tests/eval/value_cache/sparse.json.lz4
new file mode 100644
index 00000000000..0de6fae56e1
--- /dev/null
+++ b/eval/src/tests/eval/value_cache/sparse.json.lz4
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 d3ec7aae546..df0728a9d78 100644
--- a/eval/src/tests/eval/value_cache/tensor_loader_test.cpp
+++ b/eval/src/tests/eval/value_cache/tensor_loader_test.cpp
@@ -62,12 +62,20 @@ TEST_F("require that dense tensors can be loaded", ConstantTensorLoader(SimpleTe
     TEST_DO(verify_tensor(make_dense_tensor(), f1.create(TEST_PATH("dense.json"), "tensor(x[2],y[2])")));
 }
 
-TEST_F("require that sparse tensors can be loaded", ConstantTensorLoader(SimpleTensorEngine::ref())) {
-    TEST_DO(verify_tensor(make_sparse_tensor(), f1.create(TEST_PATH("sparse.json"), "tensor(x{},y{})")));
-}
-
 TEST_F("require that mixed tensors can be loaded", ConstantTensorLoader(SimpleTensorEngine::ref())) {
     TEST_DO(verify_tensor(make_mixed_tensor(), f1.create(TEST_PATH("mixed.json"), "tensor(x{},y[2])")));
 }
 
+TEST_F("require that lz4 compressed dense tensor can be loaded", ConstantTensorLoader(SimpleTensorEngine::ref())) {
+    TEST_DO(verify_tensor(make_dense_tensor(), f1.create(TEST_PATH("dense.json.lz4"), "tensor(x[2],y[2])")));
+}
+
+TEST_F("require that lz4 compressed sparse tensor can be loaded", ConstantTensorLoader(SimpleTensorEngine::ref())) {
+    TEST_DO(verify_tensor(make_sparse_tensor(), f1.create(TEST_PATH("sparse.json.lz4"), "tensor(x{},y{})")));
+}
+
+TEST_F("require that bad lz4 file fails to load creating empty result", ConstantTensorLoader(SimpleTensorEngine::ref())) {
+    TEST_DO(verify_tensor(dense_tensor_nocells(), f1.create(TEST_PATH("bad_lz4.json.lz4"), "tensor(x[2],y[2])")));
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
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 fc6181acc02..65c36ca0ef2 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
@@ -10,6 +10,7 @@
 #include <vespa/eval/eval/tensor_engine.h>
 #include <vespa/eval/eval/tensor_spec.h>
 #include <vespa/vespalib/io/mapped_file_input.h>
+#include <vespa/vespalib/data/lz4_input_decoder.h>
 
 LOG_SETUP(".vespalib.eval.value_cache.constant_tensor_loader");
 
@@ -42,6 +43,31 @@ struct AddressExtractor : ObjectTraverser {
     }
 };
 
+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());
+    }
+}
+
+void decode_json(const vespalib::string &path, Slime &slime) {
+    MappedFileInput file(path);
+    if (!file.valid()) {
+        LOG(warning, "could not read file: %s", path.c_str());
+    } else {
+        if (ends_with(path, ".lz4")) {
+            size_t buffer_size = 64 * 1024;
+            Lz4InputDecoder lz4_decoder(file, buffer_size);            
+            decode_json(path, lz4_decoder, slime);
+            if (lz4_decoder.failed()) {
+                LOG(warning, "file contains lz4 errors (%s): %s",
+                    lz4_decoder.reason().c_str(), path.c_str());
+            }
+        } else {
+            decode_json(path, file, slime);            
+        }
+    }
+}
+
 } // namespace vespalib::eval::<unnamed>
 
 using ErrorConstant = SimpleConstantValue<ErrorValue>;
@@ -57,12 +83,7 @@ ConstantTensorLoader::create(const vespalib::string &path, const vespalib::strin
         return std::make_unique<TensorConstant>(_engine.type_of(*tensor), std::move(tensor));
     }
     Slime slime;
-    MappedFileInput file(path);
-    if (!file.valid()) {
-        LOG(warning, "could not read file: %s", path.c_str());
-    } else if (slime::JsonFormat::decode(file.get(), slime) == 0) {
-        LOG(warning, "file contains invalid json: %s", path.c_str());
-    }
+    decode_json(path, slime);
     std::set<vespalib::string> indexed;
     for (const auto &dimension: value_type.dimensions()) {
         if (dimension.is_indexed()) {