Merge pull request #24292 from vespa-engine/toregge/use-array-store-allocate-in-tensor-buffer-store-store-tensor

Use ArrayStore::allocate in TensorBufferStore::store_tensor.

2 files changed, 5 insertions, 6 deletions

diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp
index 6c1b3bbd1ee..34454e9f780 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp
@@ -29,8 +29,7 @@ TensorBufferStore::TensorBufferStore(const ValueType& tensor_type, std::shared_p
       _array_store(ArrayStoreType::optimizedConfigForHugePage(max_small_subspaces_type_id,
                                                               TensorBufferTypeMapper(max_small_subspaces_type_id, &_ops),
                                                               MemoryAllocator::HUGEPAGE_SIZE, 4_Ki, 8_Ki, ALLOC_GROW_FACTOR),
-                   std::move(allocator), TensorBufferTypeMapper(max_small_subspaces_type_id, &_ops)),
-      _add_buffer()
+                   std::move(allocator), TensorBufferTypeMapper(max_small_subspaces_type_id, &_ops))
 {
 }
 
@@ -60,9 +59,10 @@ TensorBufferStore::store_tensor(const Value &tensor)
 {
     uint32_t num_subspaces = tensor.index().size();
     auto array_size = _ops.get_array_size(num_subspaces);
-    _add_buffer.resize(array_size);
-    _ops.store_tensor(_add_buffer, tensor);
-    return _array_store.add(_add_buffer);
+    auto ref = _array_store.allocate(array_size);
+    auto buf = _array_store.get_writable(ref);
+    _ops.store_tensor(buf, tensor);
+    return ref;
 }
 
 EntryRef
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h
index 14572eb07dc..18b98efa8fa 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h
@@ -23,7 +23,6 @@ class TensorBufferStore : public TensorStore
     vespalib::eval::ValueType _tensor_type;
     TensorBufferOperations    _ops;
     ArrayStoreType            _array_store;
-    std::vector<char>         _add_buffer;
 public:
     TensorBufferStore(const vespalib::eval::ValueType& tensor_type, std::shared_ptr<vespalib::alloc::MemoryAllocator> allocator, uint32_t max_small_subspaces_type_id);
     ~TensorBufferStore();