enhance: refine array view to optimize memory usage(milvus-io#38736)

Signed-off-by: MrPresent-Han <[email protected]>
MrPresent-Han · Dec 27, 2024 · 045b162 · 045b162
1 parent 18a3bc7
commit 045b162
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 14 deletions.
diff --git a/internal/core/src/common/Array.h b/internal/core/src/common/Array.h
@@ -444,13 +444,30 @@ class ArrayView {
  public:
     ArrayView() = default;
 
+    ArrayView(char* data,
+              int len,
+              size_t size,
+              DataType element_type,
+              const uint64_t* offsets_ptr):
+              data_(data),
+              length_(len),
+              size_(size),
+              element_type_(element_type),
+              offsets_ptr_(offsets_ptr){
+        AssertInfo(data!= nullptr, "data pointer for ArrayView cannot be nullptr");
+        if (IsVariableDataType(element_type_)) {
+            AssertInfo(offsets_ptr!= nullptr, "for variable data type, offsets_ptr for array view must not be nullptr");
+        }
+    }
+
     ArrayView(char* data,
               size_t size,
               DataType element_type,
               std::vector<uint64_t>&& element_offsets)
         : size_(size),
           offsets_(std::move(element_offsets)),
-          element_type_(element_type) {
+          element_type_(element_type),
+          offsets_ptr_(nullptr){
         data_ = data;
         if (IsVariableDataType(element_type_)) {
             length_ = offsets_.size();
@@ -475,10 +492,17 @@ class ArrayView {
 
         if constexpr (std::is_same_v<T, std::string> ||
                       std::is_same_v<T, std::string_view>) {
-            size_t element_length = (index == length_ - 1)
+            if (offsets_ptr_) {
+                size_t element_length = (index == length_ - 1)
+                                        ? size_ - offsets_ptr_[length_ - 1]
+                                        : offsets_ptr_[index + 1] - offsets_ptr_[index];
+                return T(data_ + offsets_ptr_[index], element_length);
+            } else {
+                size_t element_length = (index == length_ - 1)
                                         ? size_ - offsets_.back()
                                         : offsets_[index + 1] - offsets_[index];
-            return T(data_ + offsets_[index], element_length);
+                return T(data_ + offsets_[index], element_length);
+            }
         }
         if constexpr (std::is_same_v<T, int> || std::is_same_v<T, int64_t> ||
                       std::is_same_v<T, float> || std::is_same_v<T, double>) {
@@ -583,6 +607,9 @@ class ArrayView {
     // copy to result
     std::vector<uint64_t>
     get_offsets_in_copy() const {
+        if (offsets_ptr_) {
+            return std::vector<uint64_t>(offsets_ptr_, offsets_ptr_ + sizeof(uint64_t) * length_);
+        }
         return offsets_;
     }
 
@@ -663,6 +690,9 @@ class ArrayView {
     int size_ = 0;
     std::vector<uint64_t> offsets_{};
     DataType element_type_ = DataType::NONE;
+
+    //offsets for mmap
+    const uint64_t* offsets_ptr_;
 };
 
 }  // namespace milvus
diff --git a/internal/core/src/common/Chunk.cpp b/internal/core/src/common/Chunk.cpp
@@ -51,21 +51,18 @@ ArrayChunk::ConstructViews() {
         int offset = offsets_lens_[2 * i];
         int next_offset = offsets_lens_[2 * (i + 1)];
         int len = offsets_lens_[2 * i + 1];
-
         auto data_ptr = data_ + offset;
-        auto offsets_len = 0;
-        std::vector<uint64_t> element_indices = {};
+        auto offsets_bytes_len = 0;
+        uint64_t* offsets_ptr = nullptr;
         if (IsStringDataType(element_type_)) {
-            offsets_len = len * sizeof(uint64_t);
-            std::vector<uint64_t> tmp(
-                reinterpret_cast<uint64_t*>(data_ptr),
-                reinterpret_cast<uint64_t*>(data_ptr + offsets_len));
-            element_indices = std::move(tmp);
+            offsets_bytes_len = len * sizeof(uint64_t);
+            offsets_ptr = reinterpret_cast<uint64_t*>(data_ptr);
         }
-        views_.emplace_back(data_ptr + offsets_len,
-                            next_offset - offset - offsets_len,
+        views_.emplace_back(data_ptr + offsets_bytes_len,
+                            len,
+                            next_offset - offset - offsets_bytes_len,
                             element_type_,
-                            std::move(element_indices));
+                            offsets_ptr);
     }
 }