Merge pull request #14906 from vespa-engine/arnej/maybe-shrink-sparse-tensor

Arnej/maybe shrink sparse tensor

5 files changed, 32 insertions, 10 deletions

diff --git a/searchlib/src/tests/features/tensor_from_labels/tensor_from_labels_test.cpp b/searchlib/src/tests/features/tensor_from_labels/tensor_from_labels_test.cpp
index 80c98284b82..33fd7c856b3 100644
--- a/searchlib/src/tests/features/tensor_from_labels/tensor_from_labels_test.cpp
+++ b/searchlib/src/tests/features/tensor_from_labels/tensor_from_labels_test.cpp
@@ -121,8 +121,7 @@ struct ExecFixture
         test.getQueryEnv().getProperties().add("aint_query", "[11 13 17]");
     }
     const Value &extractTensor(uint32_t docid) {
-        Value::CREF value = test.resolveObjectFeature(docid);
-        return value.get();
+        return test.resolveObjectFeature(docid);
     }
     const Value &execute() {
         return extractTensor(1);
diff --git a/searchlib/src/tests/features/tensor_from_weighted_set/tensor_from_weighted_set_test.cpp b/searchlib/src/tests/features/tensor_from_weighted_set/tensor_from_weighted_set_test.cpp
index d3909223907..282c8134754 100644
--- a/searchlib/src/tests/features/tensor_from_weighted_set/tensor_from_weighted_set_test.cpp
+++ b/searchlib/src/tests/features/tensor_from_weighted_set/tensor_from_weighted_set_test.cpp
@@ -120,8 +120,7 @@ struct ExecFixture
         test.getQueryEnv().getProperties().add("wsquery", "{d:11,e:13,f:17}");
     }
     const Value &extractTensor(uint32_t docid) {
-        Value::CREF value = test.resolveObjectFeature(docid);
-        return value.get();
+        return test.resolveObjectFeature(docid);
     }
     const Value &execute() {
         return extractTensor(1);
diff --git a/searchlib/src/tests/queryeval/nearest_neighbor/nearest_neighbor_test.cpp b/searchlib/src/tests/queryeval/nearest_neighbor/nearest_neighbor_test.cpp
index 153e3cf45cc..f6ca0bd1427 100644
--- a/searchlib/src/tests/queryeval/nearest_neighbor/nearest_neighbor_test.cpp
+++ b/searchlib/src/tests/queryeval/nearest_neighbor/nearest_neighbor_test.cpp
@@ -41,8 +41,7 @@ DistanceFunction::UP euclid_d = search::tensor::make_distance_function(DistanceM
 DistanceFunction::UP euclid_f = search::tensor::make_distance_function(DistanceMetric::Euclidean, CellType::FLOAT);
 
 std::unique_ptr<Value> createTensor(const TensorSpec &spec) {
-    auto value = EngineOrFactory::get().from_spec(spec);
-    return value;
+    return EngineOrFactory::get().from_spec(spec);
 }
 
 std::unique_ptr<Value> createTensor(const vespalib::string& type_spec, double v1, double v2) {
diff --git a/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp b/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp
index 7e525097390..a6f4d56425b 100644
--- a/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp
+++ b/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp
@@ -12,13 +12,36 @@ using vespalib::datastore::EntryRef;
 using vespalib::eval::EngineOrFactory;
 using vespalib::eval::TensorSpec;
 using vespalib::eval::Value;
+using vespalib::eval::ValueType;
+using vespalib::eval::TypedCells;
+using vespalib::MemoryUsage;
 
 vespalib::string tensor_spec("tensor(x{})");
 
+class MockBigTensor : public Value
+{
+private:
+    Value::UP _real_tensor;
+public:
+    MockBigTensor(std::unique_ptr<Value> real_tensor)
+        : _real_tensor(std::move(real_tensor))
+    {}
+    MemoryUsage get_memory_usage() const override {
+        auto memuse = _real_tensor->get_memory_usage();
+        memuse.incUsedBytes(1000);
+        memuse.incAllocatedBytes(1000000);
+        return memuse;
+    }
+    const ValueType &type() const override { return _real_tensor->type(); }
+    TypedCells cells() const override { return _real_tensor->cells(); }
+    const Index &index() const override { return _real_tensor->index(); }
+};
+
 Value::UP
 make_tensor(const TensorSpec& spec)
 {
-    return EngineOrFactory::get().from_spec(spec);
+    auto value = EngineOrFactory::get().from_spec(spec);
+    return std::make_unique<MockBigTensor>(std::move(value));
 }
 
 Value::UP
@@ -59,8 +82,10 @@ TEST_F(DirectTensorStoreTest, heap_allocated_memory_is_tracked)
     auto ref = store.store_tensor(make_tensor(10));
     auto tensor_mem_usage = store.get_tensor(ref)->get_memory_usage();
     auto mem_2 = store.getMemoryUsage();
-    EXPECT_GT(tensor_mem_usage.usedBytes(), 100);
-    EXPECT_GT(tensor_mem_usage.allocatedBytes(), 500);
+    EXPECT_GT(tensor_mem_usage.usedBytes(), 500);
+    EXPECT_LT(tensor_mem_usage.usedBytes(), 50000);
+    EXPECT_GT(tensor_mem_usage.allocatedBytes(), 500000);
+    EXPECT_LT(tensor_mem_usage.allocatedBytes(), 50000000);
     EXPECT_GE(mem_2.allocatedBytes(), mem_1.allocatedBytes() + tensor_mem_usage.allocatedBytes());
     EXPECT_GT(mem_2.usedBytes(), mem_1.usedBytes() + tensor_mem_usage.allocatedBytes());
 }
diff --git a/searchlib/src/vespa/searchlib/queryeval/fake_requestcontext.h b/searchlib/src/vespa/searchlib/queryeval/fake_requestcontext.h
index 96fb2968188..c3f9685e47c 100644
--- a/searchlib/src/vespa/searchlib/queryeval/fake_requestcontext.h
+++ b/searchlib/src/vespa/searchlib/queryeval/fake_requestcontext.h
@@ -37,7 +37,7 @@ public:
         if (_query_tensor && (tensor_name == _query_tensor_name)) {
             return vespalib::eval::EngineOrFactory::get().from_spec(*_query_tensor);
         }
-        return vespalib::eval::Value::UP();
+        return {};
     }
     void set_query_tensor(const vespalib::string& name, const vespalib::eval::TensorSpec& tensor_spec) {
         _query_tensor_name = name;