Add TensorBufferStore.

13 files changed, 630 insertions, 1 deletions

diff --git a/searchlib/CMakeLists.txt b/searchlib/CMakeLists.txt
index 76dadc5605e..a7d831aa623 100644
--- a/searchlib/CMakeLists.txt
+++ b/searchlib/CMakeLists.txt
@@ -224,6 +224,7 @@ vespa_define_module(
     src/tests/tensor/hnsw_index
     src/tests/tensor/hnsw_saver
     src/tests/tensor/tensor_buffer_operations
+    src/tests/tensor/tensor_buffer_store
     src/tests/transactionlog
     src/tests/transactionlogstress
     src/tests/true
diff --git a/searchlib/src/tests/tensor/tensor_buffer_store/CMakeLists.txt b/searchlib/src/tests/tensor/tensor_buffer_store/CMakeLists.txt
new file mode 100644
index 00000000000..749d38a1383
--- /dev/null
+++ b/searchlib/src/tests/tensor/tensor_buffer_store/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(searchlib_tensor_buffer_store_test_app TEST
+    SOURCES
+    tensor_buffer_store_test.cpp
+    DEPENDS
+    searchlib
+    GTest::GTest
+)
+vespa_add_test(NAME searchlib_tensor_buffer_store_test_app COMMAND searchlib_tensor_buffer_store_test_app)
diff --git a/searchlib/src/tests/tensor/tensor_buffer_store/tensor_buffer_store_test.cpp b/searchlib/src/tests/tensor/tensor_buffer_store/tensor_buffer_store_test.cpp
new file mode 100644
index 00000000000..101b84e01aa
--- /dev/null
+++ b/searchlib/src/tests/tensor/tensor_buffer_store/tensor_buffer_store_test.cpp
@@ -0,0 +1,164 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/searchlib/tensor/tensor_buffer_store.h>
+#include <vespa/eval/eval/simple_value.h>
+#include <vespa/eval/eval/tensor_spec.h>
+#include <vespa/eval/eval/value.h>
+#include <vespa/vespalib/gtest/gtest.h>
+
+using search::tensor::TensorBufferStore;
+using vespalib::datastore::EntryRef;
+using vespalib::eval::SimpleValue;
+using vespalib::eval::TensorSpec;
+using vespalib::eval::Value;
+using vespalib::eval::ValueType;
+
+const vespalib::string tensor_type_spec("tensor(x{})");
+
+class TensorBufferStoreTest : public testing::Test
+{
+protected:
+    ValueType         _tensor_type;
+    TensorBufferStore _store;
+    TensorBufferStoreTest();
+    ~TensorBufferStoreTest() override;
+    EntryRef store_tensor(const Value& tensor);
+    EntryRef store_tensor(const TensorSpec& spec);
+    std::unique_ptr<Value> load_tensor(EntryRef ref);
+    TensorSpec load_tensor_spec(EntryRef ref);
+    vespalib::nbostream encode_stored_tensor(EntryRef ref);
+    void assert_store_load(const TensorSpec& tensor_spec);
+    void assert_store_load_many(const TensorSpec& tensor_spec);
+    void assert_store_move_load(const TensorSpec& tensor_spec);
+    void assert_store_encode_store_encoded_load(const TensorSpec& tensor_spec);
+};
+
+TensorBufferStoreTest::TensorBufferStoreTest()
+    : testing::Test(),
+      _tensor_type(ValueType::from_spec(tensor_type_spec)),
+      _store(_tensor_type, {}, 4)
+{
+}
+
+TensorBufferStoreTest::~TensorBufferStoreTest() = default;
+
+EntryRef
+TensorBufferStoreTest::store_tensor(const Value& tensor)
+{
+    EXPECT_EQ(_tensor_type, tensor.type());
+    return _store.store_tensor(tensor);
+}
+
+EntryRef
+TensorBufferStoreTest::store_tensor(const TensorSpec& spec)
+{
+    auto tensor = SimpleValue::from_spec(spec);
+    return store_tensor(*tensor);
+}
+
+std::unique_ptr<Value>
+TensorBufferStoreTest::load_tensor(EntryRef ref)
+{
+    return _store.get_tensor(ref);
+}
+
+vespalib::nbostream
+TensorBufferStoreTest::encode_stored_tensor(EntryRef ref)
+{
+    vespalib::nbostream out;
+    _store.encode_stored_tensor(ref, out);
+    return out;
+}
+
+TensorSpec
+TensorBufferStoreTest::load_tensor_spec(EntryRef ref)
+{
+    auto loaded = load_tensor(ref);
+    return TensorSpec::from_value(*loaded);
+}
+
+void
+TensorBufferStoreTest::assert_store_load(const TensorSpec& tensor_spec)
+{
+    auto ref = store_tensor(tensor_spec);
+    auto loaded_spec = load_tensor_spec(ref);
+    _store.holdTensor(ref);
+    EXPECT_EQ(tensor_spec, loaded_spec);
+}
+
+void
+TensorBufferStoreTest::assert_store_load_many(const TensorSpec& tensor_spec)
+{
+    constexpr uint32_t cnt = 2000;
+    std::vector<EntryRef> refs;
+    for (uint32_t i = 0; i < cnt; ++i) {
+        refs.emplace_back(store_tensor(tensor_spec));
+    }
+    for (auto ref : refs) {
+        auto loaded_spec = load_tensor_spec(ref);
+        _store.holdTensor(ref);
+        EXPECT_EQ(tensor_spec, loaded_spec);
+    }
+}
+
+void
+TensorBufferStoreTest::assert_store_move_load(const TensorSpec& tensor_spec)
+{
+    auto ref = store_tensor(tensor_spec);
+    auto ref2 = _store.move(ref);
+    EXPECT_NE(ref, ref2);
+    auto loaded_spec = load_tensor_spec(ref2);
+    _store.holdTensor(ref2);
+    EXPECT_EQ(tensor_spec, loaded_spec);
+}
+
+void
+TensorBufferStoreTest::assert_store_encode_store_encoded_load(const TensorSpec& tensor_spec)
+{
+    auto ref = store_tensor(tensor_spec);
+    auto encoded = encode_stored_tensor(ref);
+    _store.holdTensor(ref);
+    auto ref2 = _store.store_encoded_tensor(encoded);
+    EXPECT_NE(ref, ref2);
+    auto loaded_spec = load_tensor_spec(ref2);
+    _store.holdTensor(ref2);
+    EXPECT_EQ(tensor_spec, loaded_spec);
+}
+
+std::vector<TensorSpec> tensor_specs = {
+    TensorSpec(tensor_type_spec),
+    TensorSpec(tensor_type_spec).add({{"x", "a"}}, 4.5),
+    TensorSpec(tensor_type_spec).add({{"x", "a"}}, 4.5).add({{"x", "b"}}, 5.5),
+    TensorSpec(tensor_type_spec).add({{"x", "a"}}, 4.5).add({{"x", "b"}}, 5.5).add({{"x", "c"}}, 6.5),
+    TensorSpec(tensor_type_spec).add({{"x", "a"}}, 4.5).add({{"x", "b"}}, 5.5).add({{"x", "c"}}, 6.5).add({{"x", "d"}}, 7.5)
+};
+
+TEST_F(TensorBufferStoreTest, tensor_can_be_stored_and_loaded)
+{
+    for (auto& tensor_spec : tensor_specs) {
+        assert_store_load(tensor_spec);
+    }
+}
+
+TEST_F(TensorBufferStoreTest, tensor_can_be_stored_and_loaded_many_times)
+{
+    for (auto& tensor_spec : tensor_specs) {
+        assert_store_load_many(tensor_spec);
+    }
+}
+
+TEST_F(TensorBufferStoreTest, stored_tensor_can_be_copied)
+{
+    for (auto& tensor_spec : tensor_specs) {
+        assert_store_move_load(tensor_spec);
+    }
+}
+
+TEST_F(TensorBufferStoreTest, stored_tensor_can_be_encoded_and_stored_as_encoded_and_loaded)
+{
+    for (auto& tensor_spec : tensor_specs) {
+        assert_store_encode_store_encoded_load(tensor_spec);
+    }
+}
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt b/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt
index 7815ef7e770..46bfc0909aa 100644
--- a/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt
+++ b/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt
@@ -24,13 +24,17 @@ vespa_add_library(searchlib_tensor OBJECT
     imported_tensor_attribute_vector_read_guard.cpp
     inner_product_distance.cpp
     inv_log_level_generator.cpp
+    large_subspaces_buffer_type.cpp
     nearest_neighbor_index.cpp
     nearest_neighbor_index_saver.cpp
     serialized_fast_value_attribute.cpp
+    small_subspaces_buffer_type.cpp
     streamed_value_saver.cpp
     streamed_value_store.cpp
     tensor_attribute.cpp
     tensor_buffer_operations.cpp
+    tensor_buffer_store.cpp
+    tensor_buffer_type_mapper.cpp
     tensor_deserialize.cpp
     tensor_store.cpp
     reusable_set_visited_tracker.cpp
diff --git a/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.cpp b/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.cpp
new file mode 100644
index 00000000000..cdd4d35c1df
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.cpp
@@ -0,0 +1,86 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "large_subspaces_buffer_type.h"
+#include "tensor_buffer_operations.h"
+#include "tensor_buffer_type_mapper.h"
+#include <vespa/vespalib/datastore/buffer_type.hpp>
+#include <vespa/vespalib/util/array.hpp>
+#include <vespa/vespalib/util/arrayref.h>
+
+using vespalib::alloc::MemoryAllocator;
+
+namespace search::tensor {
+
+LargeSubspacesBufferType::LargeSubspacesBufferType(const AllocSpec& spec, std::shared_ptr<MemoryAllocator> memory_allocator, TensorBufferTypeMapper& type_mapper) noexcept
+    : ParentType(1u, spec.minArraysInBuffer, spec.maxArraysInBuffer, spec.numArraysForNewBuffer, spec.allocGrowFactor),
+      _memory_allocator(std::move(memory_allocator)),
+      _ops(type_mapper.get_tensor_buffer_operations())
+{
+}
+
+LargeSubspacesBufferType::~LargeSubspacesBufferType() = default;
+
+void
+LargeSubspacesBufferType::cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx)
+{
+    auto elem = static_cast<ArrayType*>(buffer) + offset;
+    for (size_t i = 0; i < numElems; ++i) {
+        if (!elem->empty()) {
+            cleanCtx.extraBytesCleaned(elem->size());
+            _ops.reclaim_labels({elem->data(), elem->size()});
+            ArrayType().swap(*elem);
+        }
+        ++elem;
+    }
+}
+
+void
+LargeSubspacesBufferType::destroyElements(void *buffer, ElemCount numElems)
+{
+    auto elem = static_cast<ArrayType*>(buffer);
+    for (size_t i = 0; i < numElems; ++i) {
+        if (!elem->empty()) {
+            _ops.reclaim_labels({elem->data(), elem->size()});
+            ArrayType().swap(*elem);
+        }
+        ++elem;
+    }
+}
+
+void
+LargeSubspacesBufferType::fallbackCopy(void *newBuffer, const void *oldBuffer, ElemCount numElems)
+{
+    auto old_elems = static_cast<const ArrayType*>(oldBuffer);
+    auto new_elems = static_cast<ArrayType*>(newBuffer);
+    for (size_t i = 0; i < numElems; ++i) {
+        auto& old_elem = old_elems[i];
+        new (new_elems + i) ArrayType(old_elem);
+        if (!old_elem.empty()) {
+            _ops.copied_labels({old_elem.data(), old_elem.size()});
+        }
+    }
+}
+
+void
+LargeSubspacesBufferType::initializeReservedElements(void *buffer, ElemCount reservedElements)
+{
+    auto new_elems = static_cast<ArrayType*>(buffer);
+    const auto& empty = empty_entry();
+    for (size_t i = 0; i < reservedElements; ++i) {
+        new (new_elems + i) ArrayType(empty);
+    }
+}
+
+const vespalib::alloc::MemoryAllocator*
+LargeSubspacesBufferType::get_memory_allocator() const
+{
+    return _memory_allocator.get();
+}
+
+}
+
+namespace vespalib::datastore {
+
+template class BufferType<Array<char>>;
+
+}
diff --git a/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.h b/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.h
new file mode 100644
index 00000000000..cfab8ef20af
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.h
@@ -0,0 +1,40 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/vespalib/datastore/array_store_config.h>
+#include <vespa/vespalib/datastore/buffer_type.h>
+#include <vespa/vespalib/util/array.h>
+#include <memory>
+
+namespace vespalib::alloc { class MemoryAllocator; }
+
+namespace search::tensor {
+
+class TensorBufferOperations;
+class TensorBufferTypeMapper;
+
+/*
+ * Class representing buffer type for tensors with a large number of
+ * subspaces in array store. Tensor buffers are externally allocated
+ * (cf. vespalib::Array).
+ */
+class LargeSubspacesBufferType : public vespalib::datastore::BufferType<vespalib::Array<char>>
+{
+    using AllocSpec = vespalib::datastore::ArrayStoreConfig::AllocSpec;
+    using ArrayType = vespalib::Array<char>;
+    using ParentType = vespalib::datastore::BufferType<ArrayType>;
+    using CleanContext = typename ParentType::CleanContext;
+    std::shared_ptr<vespalib::alloc::MemoryAllocator> _memory_allocator;
+    TensorBufferOperations& _ops;
+public:
+    LargeSubspacesBufferType(const AllocSpec& spec, std::shared_ptr<vespalib::alloc::MemoryAllocator> memory_allocator, TensorBufferTypeMapper& type_mapper) noexcept;
+    ~LargeSubspacesBufferType() override;
+    void cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx) override;
+    void destroyElements(void *buffer, ElemCount numElems) override;
+    void fallbackCopy(void *newBuffer, const void *oldBuffer, ElemCount numElems) override;
+    void initializeReservedElements(void *buffer, ElemCount reservedElements) override;
+    const vespalib::alloc::MemoryAllocator* get_memory_allocator() const override;
+};
+
+}
diff --git a/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.cpp b/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.cpp
new file mode 100644
index 00000000000..adbd3dee2b7
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.cpp
@@ -0,0 +1,67 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "small_subspaces_buffer_type.h"
+#include "tensor_buffer_operations.h"
+#include "tensor_buffer_type_mapper.h"
+#include <vespa/vespalib/util/arrayref.h>
+
+using vespalib::alloc::MemoryAllocator;
+
+namespace search::tensor {
+
+SmallSubspacesBufferType::SmallSubspacesBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr<MemoryAllocator> memory_allocator, TensorBufferTypeMapper& type_mapper) noexcept
+    : ParentType(array_size, spec.minArraysInBuffer, spec.maxArraysInBuffer, spec.numArraysForNewBuffer, spec.allocGrowFactor),
+      _memory_allocator(std::move(memory_allocator)),
+      _ops(type_mapper.get_tensor_buffer_operations())
+{
+}
+
+SmallSubspacesBufferType::~SmallSubspacesBufferType() = default;
+
+void
+SmallSubspacesBufferType::cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext)
+{
+    char* elem = static_cast<char *>(buffer) + offset;
+    while (numElems >= getArraySize()) {
+        _ops.reclaim_labels(vespalib::ArrayRef<char>(elem, getArraySize()));
+        elem += getArraySize();
+        numElems -= getArraySize();
+    }
+}
+
+void
+SmallSubspacesBufferType::destroyElements(void *buffer, ElemCount numElems)
+{
+    char* elem = static_cast<char *>(buffer);
+    while (numElems >= getArraySize()) {
+        _ops.reclaim_labels(vespalib::ArrayRef<char>(elem, getArraySize()));
+        elem += getArraySize();
+        numElems -= getArraySize();
+    }
+}
+
+void
+SmallSubspacesBufferType::fallbackCopy(void *newBuffer, const void *oldBuffer, ElemCount numElems)
+{
+    memcpy(newBuffer, oldBuffer, numElems);
+    const char *elem = static_cast<const char *>(oldBuffer);
+    while (numElems >= getArraySize()) {
+        _ops.copied_labels(vespalib::ConstArrayRef<char>(elem, getArraySize()));
+        elem += getArraySize();
+        numElems -= getArraySize();
+    }
+}
+
+void
+SmallSubspacesBufferType::initializeReservedElements(void *buffer, ElemCount reservedElements)
+{
+    memset(buffer, 0, reservedElements);
+}
+
+const vespalib::alloc::MemoryAllocator*
+SmallSubspacesBufferType::get_memory_allocator() const
+{
+    return _memory_allocator.get();
+}
+
+}
diff --git a/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.h b/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.h
new file mode 100644
index 00000000000..a778183c5a2
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.h
@@ -0,0 +1,40 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/vespalib/datastore/array_store_config.h>
+#include <vespa/vespalib/datastore/buffer_type.h>
+#include <memory>
+
+namespace vespalib::alloc { class MemoryAllocator; }
+
+namespace search::tensor {
+
+class TensorBufferOperations;
+class TensorBufferTypeMapper;
+
+/*
+ * Class representing buffer type for tensors with a small number of
+ * subspaces in array store. Tensor buffers are internal in data store buffer.
+ */
+class SmallSubspacesBufferType : public vespalib::datastore::BufferType<char>
+{
+    using AllocSpec = vespalib::datastore::ArrayStoreConfig::AllocSpec;
+    using ParentType = vespalib::datastore::BufferType<char>;
+    std::shared_ptr<vespalib::alloc::MemoryAllocator> _memory_allocator;
+    TensorBufferOperations& _ops;
+public:
+    SmallSubspacesBufferType(const SmallSubspacesBufferType&) = delete;
+    SmallSubspacesBufferType& operator=(const SmallSubspacesBufferType&) = delete;
+    SmallSubspacesBufferType(SmallSubspacesBufferType&&) noexcept = default;
+    SmallSubspacesBufferType& operator=(SmallSubspacesBufferType&&) noexcept = default;
+    SmallSubspacesBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr<vespalib::alloc::MemoryAllocator> memory_allocator, TensorBufferTypeMapper& type_mapper) noexcept;
+    ~SmallSubspacesBufferType() override;
+    void cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx) override;
+    void destroyElements(void *buffer, ElemCount numElems) override;
+    void fallbackCopy(void *newBuffer, const void *oldBuffer, ElemCount numElems) override;
+    void initializeReservedElements(void *buffer, ElemCount reservedElements) override;
+    const vespalib::alloc::MemoryAllocator* get_memory_allocator() const override;
+};
+
+}
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp
new file mode 100644
index 00000000000..6c1b3bbd1ee
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp
@@ -0,0 +1,97 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "tensor_buffer_store.h"
+#include <vespa/eval/eval/value_codec.h>
+#include <vespa/eval/streamed/streamed_value_builder_factory.h>
+#include <vespa/vespalib/datastore/array_store.hpp>
+#include <vespa/vespalib/datastore/buffer_type.hpp>
+#include <vespa/vespalib/datastore/datastore.hpp>
+#include <vespa/vespalib/util/size_literals.h>
+
+using vespalib::alloc::MemoryAllocator;
+using vespalib::datastore::EntryRef;
+using vespalib::eval::StreamedValueBuilderFactory;
+using vespalib::eval::Value;
+using vespalib::eval::ValueType;
+
+namespace search::tensor {
+
+namespace {
+
+constexpr float ALLOC_GROW_FACTOR = 0.2;
+
+}
+
+TensorBufferStore::TensorBufferStore(const ValueType& tensor_type, std::shared_ptr<MemoryAllocator> allocator, uint32_t max_small_subspaces_type_id)
+    : TensorStore(ArrayStoreType::get_data_store_base(_array_store)),
+      _tensor_type(tensor_type),
+      _ops(_tensor_type),
+      _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()
+{
+}
+
+TensorBufferStore::~TensorBufferStore() = default;
+
+void
+TensorBufferStore::holdTensor(EntryRef ref)
+{
+    _array_store.remove(ref);
+}
+
+EntryRef
+TensorBufferStore::move(EntryRef ref)
+{
+    if (!ref.valid()) {
+        return EntryRef();
+    }
+    auto buf = _array_store.get(ref);
+    auto new_ref = _array_store.add(buf);
+    _ops.copied_labels(buf);
+    _array_store.remove(ref);
+    return new_ref;
+}
+
+EntryRef
+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);
+}
+
+EntryRef
+TensorBufferStore::store_encoded_tensor(vespalib::nbostream &encoded)
+{
+    const auto &factory = StreamedValueBuilderFactory::get();
+    auto val = vespalib::eval::decode_value(encoded, factory);
+    return store_tensor(*val);
+}
+
+std::unique_ptr<Value>
+TensorBufferStore::get_tensor(EntryRef ref) const
+{
+    if (!ref.valid()) {
+        return {};
+    }
+    auto buf = _array_store.get(ref);
+    return _ops.make_fast_view(buf, _tensor_type);
+}
+
+bool
+TensorBufferStore::encode_stored_tensor(EntryRef ref, vespalib::nbostream &target) const
+{
+    if (!ref.valid()) {
+        return false;
+    }
+    auto buf = _array_store.get(ref);
+    _ops.encode_stored_tensor(buf, _tensor_type, target);
+    return true;
+}
+
+}
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h
new file mode 100644
index 00000000000..14572eb07dc
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h
@@ -0,0 +1,38 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "tensor_store.h"
+#include "tensor_buffer_operations.h"
+#include "tensor_buffer_type_mapper.h"
+#include "large_subspaces_buffer_type.h"
+#include "small_subspaces_buffer_type.h"
+#include <vespa/eval/eval/value_type.h>
+#include <vespa/vespalib/datastore/array_store.h>
+
+namespace search::tensor {
+
+/**
+ * Class for storing tensor buffers in memory and making tensor views
+ * based on stored tensor buffer.
+ */
+class TensorBufferStore : public TensorStore
+{
+    using RefType = vespalib::datastore::EntryRefT<19>;
+    using ArrayStoreType = vespalib::datastore::ArrayStore<char, RefType, TensorBufferTypeMapper>;
+    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();
+    void holdTensor(EntryRef ref) override;
+    EntryRef move(EntryRef ref) override;
+    EntryRef store_tensor(const vespalib::eval::Value &tensor);
+    EntryRef store_encoded_tensor(vespalib::nbostream &encoded);
+    std::unique_ptr<vespalib::eval::Value> get_tensor(EntryRef ref) const;
+    bool encode_stored_tensor(EntryRef ref, vespalib::nbostream &target) const;
+};
+
+}
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.cpp
new file mode 100644
index 00000000000..b4b0b9bbc79
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.cpp
@@ -0,0 +1,47 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "tensor_buffer_type_mapper.h"
+#include "tensor_buffer_operations.h"
+#include <algorithm>
+
+namespace search::tensor {
+
+TensorBufferTypeMapper::TensorBufferTypeMapper()
+    : _array_sizes(),
+      _ops(nullptr)
+{
+}
+
+TensorBufferTypeMapper::TensorBufferTypeMapper(uint32_t max_small_subspaces_type_id, TensorBufferOperations* ops)
+    : _array_sizes(),
+      _ops(ops)
+{
+    _array_sizes.reserve(max_small_subspaces_type_id + 1);
+    _array_sizes.emplace_back(0); // type id 0 uses LargeSubspacesBufferType
+    for (uint32_t type_id = 1; type_id <= max_small_subspaces_type_id; ++type_id) {
+        auto num_subspaces = type_id - 1;
+        _array_sizes.emplace_back(_ops->get_array_size(num_subspaces));
+    }
+}
+
+TensorBufferTypeMapper::~TensorBufferTypeMapper() = default;
+
+uint32_t
+TensorBufferTypeMapper::get_type_id(size_t array_size) const
+{
+    assert(!_array_sizes.empty());
+    auto result = std::lower_bound(_array_sizes.begin() + 1, _array_sizes.end(), array_size);
+    if (result == _array_sizes.end()) {
+        return 0; // type id 0 uses LargeSubspacesBufferType
+    }
+    return result - _array_sizes.begin();
+}
+
+size_t
+TensorBufferTypeMapper::get_array_size(uint32_t type_id) const
+{
+    assert(type_id > 0 && type_id < _array_sizes.size());
+    return _array_sizes[type_id];
+}
+
+}
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.h b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.h
new file mode 100644
index 00000000000..1e02c1cb608
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.h
@@ -0,0 +1,35 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+namespace search::tensor {
+
+class LargeSubspacesBufferType;
+class SmallSubspacesBufferType;
+class TensorBufferOperations;
+
+/*
+ * This class provides mapping between type ids and array sizes needed for
+ * storing a tensor.
+ */
+class TensorBufferTypeMapper
+{
+    std::vector<size_t> _array_sizes;
+    TensorBufferOperations* _ops;
+public:
+    using SmallBufferType = SmallSubspacesBufferType;
+    using LargeBufferType = LargeSubspacesBufferType;
+
+    TensorBufferTypeMapper();
+    TensorBufferTypeMapper(uint32_t max_small_subspaces_type_id, TensorBufferOperations* ops);
+    ~TensorBufferTypeMapper();
+
+    uint32_t get_type_id(size_t array_size) const;
+    size_t get_array_size(uint32_t type_id) const;
+    TensorBufferOperations& get_tensor_buffer_operations() const noexcept { return *_ops; }
+};
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/array_store.h b/vespalib/src/vespa/vespalib/datastore/array_store.h
index fd0271915dd..2c6aaea6b61 100644
--- a/vespalib/src/vespa/vespalib/datastore/array_store.h
+++ b/vespalib/src/vespa/vespalib/datastore/array_store.h
@@ -110,7 +110,8 @@ public:
     static vespalib::GenerationHolder &getGenerationHolderLocation(ArrayStore &self) {
         return DataStoreBase::getGenerationHolderLocation(self._store);
     }
-
+    // need object location before construction
+    static DataStoreBase& get_data_store_base(ArrayStore &self) { return self._store; }
 
     // Should only be used for unit testing
     const BufferState &bufferState(EntryRef ref) const;