Merge pull request #26335 from vespa-engine/toregge/split-tensor-buffer-type-mapper

Split tensor buffer type mapper.

9 files changed, 104 insertions, 25 deletions

diff --git a/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt b/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt
index f11004363f8..1c3b9112dda 100644
--- a/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt
+++ b/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt
@@ -3,6 +3,7 @@ vespa_add_library(vespalib_vespalib_datastore OBJECT
     SOURCES
     array_store.cpp
     array_store_config.cpp
+    array_store_type_mapper.cpp
     atomic_entry_ref.cpp
     buffer_free_list.cpp
     buffer_stats.cpp
diff --git a/vespalib/src/vespa/vespalib/datastore/array_store.h b/vespalib/src/vespa/vespalib/datastore/array_store.h
index 9e403a52b44..2a193172959 100644
--- a/vespalib/src/vespa/vespalib/datastore/array_store.h
+++ b/vespalib/src/vespa/vespalib/datastore/array_store.h
@@ -2,7 +2,7 @@
 
 #pragma once
 
-#include "array_store_type_mapper.h"
+#include "array_store_simple_type_mapper.h"
 #include "array_store_config.h"
 #include "buffer_type.h"
 #include "bufferstate.h"
@@ -29,7 +29,7 @@ namespace vespalib::datastore {
  *
  * The max value of maxSmallArrayTypeId is (2^bufferBits - 1).
  */
-template <typename EntryT, typename RefT = EntryRefT<19>, typename TypeMapperT = ArrayStoreTypeMapper<EntryT> >
+template <typename EntryT, typename RefT = EntryRefT<19>, typename TypeMapperT = ArrayStoreSimpleTypeMapper<EntryT> >
 class ArrayStore : public ICompactable
 {
 public:
diff --git a/vespalib/src/vespa/vespalib/datastore/array_store_simple_type_mapper.h b/vespalib/src/vespa/vespalib/datastore/array_store_simple_type_mapper.h
new file mode 100644
index 00000000000..a0cd7827b2d
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/array_store_simple_type_mapper.h
@@ -0,0 +1,29 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "large_array_buffer_type.h"
+#include "small_array_buffer_type.h"
+
+namespace vespalib::datastore {
+
+/**
+ * This class provides a 1-to-1 mapping between type ids and array sizes for small arrays in an array store.
+ *
+ * This means that buffers for type id 1 stores arrays of size 1, buffers for type id 2 stores arrays of size 2, and so on.
+ * Type id 0 is always reserved for large arrays allocated on the heap.
+ *
+ * A more complex mapping can be used by creating a custom mapper and BufferType implementations.
+ */
+template <typename EntryT>
+class ArrayStoreSimpleTypeMapper {
+public:
+    using SmallBufferType = SmallArrayBufferType<EntryT>;
+    using LargeBufferType = LargeArrayBufferType<EntryT>;
+
+    uint32_t get_type_id(size_t array_size) const { return array_size; }
+    size_t get_array_size(uint32_t type_id) const { return type_id; }
+    static uint32_t get_max_small_array_type_id(uint32_t max_small_array_type_id) noexcept { return max_small_array_type_id; }
+};
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/array_store_type_mapper.cpp b/vespalib/src/vespa/vespalib/datastore/array_store_type_mapper.cpp
new file mode 100644
index 00000000000..ff514f5a00b
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/array_store_type_mapper.cpp
@@ -0,0 +1,41 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "array_store_type_mapper.h"
+#include <algorithm>
+#include <cassert>
+
+namespace vespalib::datastore {
+
+ArrayStoreTypeMapper::ArrayStoreTypeMapper()
+    : _array_sizes()
+{
+}
+
+ArrayStoreTypeMapper::~ArrayStoreTypeMapper() = default;
+
+uint32_t
+ArrayStoreTypeMapper::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 buffer type for large arrays
+    }
+    return result - _array_sizes.begin();
+}
+
+size_t
+ArrayStoreTypeMapper::get_array_size(uint32_t type_id) const
+{
+    assert(type_id > 0 && type_id < _array_sizes.size());
+    return _array_sizes[type_id];
+}
+
+uint32_t
+ArrayStoreTypeMapper::get_max_small_array_type_id(uint32_t max_small_array_type_id) const noexcept
+{
+    auto clamp_type_id = _array_sizes.size() - 1;
+    return (clamp_type_id < max_small_array_type_id) ? clamp_type_id : max_small_array_type_id;
+}
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/array_store_type_mapper.h b/vespalib/src/vespa/vespalib/datastore/array_store_type_mapper.h
index 1b1e085ed16..e707627de19 100644
--- a/vespalib/src/vespa/vespalib/datastore/array_store_type_mapper.h
+++ b/vespalib/src/vespa/vespalib/datastore/array_store_type_mapper.h
@@ -2,28 +2,30 @@
 
 #pragma once
 
-#include "large_array_buffer_type.h"
-#include "small_array_buffer_type.h"
+#include <cstddef>
+#include <cstdint>
+#include <vector>
 
 namespace vespalib::datastore {
 
-/**
- * This class provides a 1-to-1 mapping between type ids and array sizes for small arrays in an array store.
+/*
+ * This class provides mapping between type ids and array sizes needed for
+ * storing a value with size smaller than or equal to the array size.
  *
- * This means that buffers for type id 1 stores arrays of size 1, buffers for type id 2 stores arrays of size 2, and so on.
- * Type id 0 is always reserved for large arrays allocated on the heap.
- *
- * A more complex mapping can be used by creating a custom mapper and BufferType implementations.
+ * The array sizes vector is a monotic increasing sequence that might end
+ * with exponential growth.
  */
-template <typename EntryT>
-class ArrayStoreTypeMapper {
+class ArrayStoreTypeMapper
+{
+protected:
+    std::vector<size_t> _array_sizes;
 public:
-    using SmallBufferType = SmallArrayBufferType<EntryT>;
-    using LargeBufferType = LargeArrayBufferType<EntryT>;
+    ArrayStoreTypeMapper();
+    ~ArrayStoreTypeMapper();
 
-    uint32_t get_type_id(size_t array_size) const { return array_size; }
-    size_t get_array_size(uint32_t type_id) const { return type_id; }
-    static uint32_t get_max_small_array_type_id(uint32_t max_small_array_type_id) noexcept { return max_small_array_type_id; }
+    uint32_t get_type_id(size_t array_size) const;
+    size_t get_array_size(uint32_t type_id) const;
+    uint32_t get_max_small_array_type_id(uint32_t max_small_array_type_id) const noexcept;
 };
 
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h
index 87083b2a878..600925969a3 100644
--- a/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h
+++ b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h
@@ -11,8 +11,6 @@ namespace vespalib::alloc { class MemoryAllocator; }
 
 namespace vespalib::datastore {
 
-template <typename EntryT> class ArrayStoreTypeMapper;
-
 /*
  * Class representing buffer type for large arrays in ArrayStore
  */
@@ -26,7 +24,12 @@ class LargeArrayBufferType : public BufferType<Array<EntryT>>
     using CleanContext = typename ParentType::CleanContext;
     std::shared_ptr<alloc::MemoryAllocator> _memory_allocator;
 public:
-    LargeArrayBufferType(const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator, ArrayStoreTypeMapper<EntryT>& mapper) noexcept;
+    LargeArrayBufferType(const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator) noexcept;
+    template <typename TypeMapper>
+    LargeArrayBufferType(const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator, TypeMapper&) noexcept
+        : LargeArrayBufferType(spec, std::move(memory_allocator))
+    {
+    }
     ~LargeArrayBufferType() override;
     void cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx) override;
     const vespalib::alloc::MemoryAllocator* get_memory_allocator() const override;
diff --git a/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.hpp b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.hpp
index f4d1d1c90a1..3042bbff73f 100644
--- a/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.hpp
@@ -8,7 +8,7 @@
 namespace vespalib::datastore {
 
 template <typename EntryT>
-LargeArrayBufferType<EntryT>::LargeArrayBufferType(const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator, ArrayStoreTypeMapper<EntryT>&) noexcept
+LargeArrayBufferType<EntryT>::LargeArrayBufferType(const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator) noexcept
     : BufferType<Array<EntryT>>(1u, spec.minArraysInBuffer, spec.maxArraysInBuffer, spec.numArraysForNewBuffer, spec.allocGrowFactor),
       _memory_allocator(std::move(memory_allocator))
 {
diff --git a/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.h b/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.h
index f6cb860348d..676a9d3790f 100644
--- a/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.h
+++ b/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.h
@@ -10,8 +10,6 @@ namespace vespalib::alloc { class MemoryAllocator; }
 
 namespace vespalib::datastore {
 
-template <typename EntryT> class ArrayStoreTypeMapper;
-
 /*
  * Class representing buffer type for small arrays in ArrayStore
  */
@@ -25,7 +23,12 @@ public:
     SmallArrayBufferType& operator=(const SmallArrayBufferType&) = delete;
     SmallArrayBufferType(SmallArrayBufferType&&) noexcept = default;
     SmallArrayBufferType& operator=(SmallArrayBufferType&&) noexcept = default;
-    SmallArrayBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator, ArrayStoreTypeMapper<EntryT>& mapper) noexcept;
+    SmallArrayBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator) noexcept;
+    template <typename TypeMapper>
+    SmallArrayBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator, TypeMapper&) noexcept
+        : SmallArrayBufferType(array_size, spec, std::move(memory_allocator))
+    {
+    }
     ~SmallArrayBufferType() override;
     const vespalib::alloc::MemoryAllocator* get_memory_allocator() const override;
 };
diff --git a/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.hpp b/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.hpp
index 9068b8db7b1..414804417eb 100644
--- a/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.hpp
@@ -7,7 +7,7 @@
 namespace vespalib::datastore {
 
 template <typename EntryT>
-SmallArrayBufferType<EntryT>::SmallArrayBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator, ArrayStoreTypeMapper<EntryT>&) noexcept
+SmallArrayBufferType<EntryT>::SmallArrayBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator) noexcept
     : BufferType<EntryT>(array_size, spec.minArraysInBuffer, spec.maxArraysInBuffer, spec.numArraysForNewBuffer, spec.allocGrowFactor),
       _memory_allocator(std::move(memory_allocator))
 {