Repository: antgroup/vsag (language: python)

Issue:
Add a new Layer "DataCell" to manage data access 
I think high level index should ignore the specific implement of quantization and io type, So the quantization and io components can be managed by new layer DataCell which provide different data access interface to the index layer. just like the following show

|Layer| Description|Implement|
|-------|--------|---------|
|Index Layer | ANN algorithm & architecture.          |                       Graph (HNSW,NNDescent) + Reorder,  IVF(PQ,SQ,fastscan), HNSW_IVF, MultiIndex ......|
|DataCell Layer| data access, use quantization, io and others    | FlattenDC, GraphNeighborhoodDC, MixtureDC .......|
|Operator Layer| quantizations, IO and other basic components                                                 | PQ,SQ,MemoryIO,SSDIO,......|

this abstract layered approach may facilitate scalability and code reuse
currently we only have index layer, so the issue will provide Datacell Layer and Operator Layer in this feature.
the related pull requests:

- #25 
- #34 
- #37 
- #41 
- #35 
- #59 
- #60 
- #61 
- #62 
- #79 
- #80 
- #96 
- #111 
- #114 
- #125 
todo tasks:
- support mixture DataCell
- adapt Graph DataCell for HNSW
- implement HNSW + SQ4 + SQ8 Index 




Discussion / hints:

Fix patch:
diff --git a/src/data_cell/flatten_datacell.h b/src/data_cell/flatten_datacell.h
index cde92da8..6b388fef 100644
--- a/src/data_cell/flatten_datacell.h
+++ b/src/data_cell/flatten_datacell.h
@@ -67,7 +67,10 @@ class FlattenDataCell : public FlattenInterface {
     }
 
     [[nodiscard]] const uint8_t*
-    GetCodesById(uint64_t id) const override;
+    GetCodesById(uint64_t id, bool& need_release) const override;
+
+    bool
+    GetCodesById(uint64_t id, uint8_t* codes) const override;
 
     inline void
     SetQuantizer(std::shared_ptr<Quantizer<QuantTmpl>> quantizer) {
@@ -165,23 +168,41 @@ FlattenDataCell<QuantTmpl, IOTmpl>::query(float* result_dists,
                                           const uint64_t* idx,
                                           uint64_t id_count) {
     for (uint64_t i = 0; i < id_count; ++i) {
-        const auto* codes = this->GetCodesById(idx[i]);
+        bool release = false;
+        const auto* codes = this->GetCodesById(idx[i], release);
         computer->ComputeDist(codes, result_dists + i);
+        if (release) {
+            this->io_->Release(codes);
+        }
     }
 }
 
 template <typename QuantTmpl, typename IOTmpl>
 float
 FlattenDataCell<QuantTmpl, IOTmpl>::ComputePairVectors(uint64_t id1, uint64_t id2) {
-    const auto* codes1 = this->GetCodesById(id1);
-    const auto* codes2 = this->GetCodesById(id1);
-    return this->quantizer_->Compute(codes1, codes2);
+    bool release1, release2;
+    const auto* codes1 = this->GetCodesById(id1, release1);
+    const auto* codes2 = this->GetCodesById(id1, release2);
+    auto result = this->quantizer_->Compute(codes1, codes2);
+    if (release1) {
+        io_->Release(codes1);
+    }
+    if (release2) {
+        io_->Release(codes2);
+    }
+
+    return result;
 }
 
 template <typename QuantTmpl, typename IOTmpl>
 const uint8_t*
-FlattenDataCell<QuantTmpl, IOTmpl>::GetCodesById(uint64_t id) const {
-    return io_->Read(code_size_, id * code_size_);
+FlattenDataCell<QuantTmpl, IOTmpl>::GetCodesById(uint64_t id, bool& need_release) const {
+    return io_->Read(code_size_, id * code_size_, need_release);
 }
 
+template <typename QuantTmpl, typename IOTmpl>
+bool
+FlattenDataCell<QuantTmpl, IOTmpl>::GetCodesById(uint64_t id, uint8_t* codes) const {
+    return io_->Read(code_size_, id * code_size_, codes);
+}
 }  // namespace vsag
diff --git a/src/data_cell/flatten_interface.h b/src/data_cell/flatten_interface.h
index daa01097..e2027ce6 100644
--- a/src/data_cell/flatten_interface.h
+++ b/src/data_cell/flatten_interface.h
@@ -48,10 +48,15 @@ class FlattenInterface {
     };
 
     [[nodiscard]] virtual const uint8_t*
-    GetCodesById(uint64_t id) const {
+    GetCodesById(uint64_t id, bool& need_release) const {
         return nullptr;
     }
 
+    virtual bool
+    GetCodesById(uint64_t id, uint8_t* codes) const {
+        return false;
+    }
+
     [[nodiscard]] virtual uint64_t
     TotalCount() const {
         return this->total_count_;
diff --git a/src/io/basic_io.h b/src/io/basic_io.h
index df501da8..afaa1b97 100644
--- a/src/io/basic_io.h
+++ b/src/io/basic_io.h
@@ -35,13 +35,13 @@ class BasicIO {
     }
 
     inline bool
-    Read(uint8_t* data, uint64_t size, uint64_t offset) const {
-        return cast().ReadImpl(data, size, offset);
+    Read(uint64_t size, uint64_t offset, uint8_t* data) const {
+        return cast().ReadImpl(size, offset, data);
     }
 
     [[nodiscard]] inline const uint8_t*
-    Read(uint64_t size, uint64_t offset) const {
-        return cast().ReadImpl(size, offset);
+    Read(uint64_t size, uint64_t offset, bool& need_release) const {
+        return cast().ReadImpl(size, offset, need_release);  // TODO(LHT129): use IOReadObject
     }
 
     inline bool
@@ -64,6 +64,11 @@ class BasicIO {
         return cast().DeserializeImpl(reader);
     }
 
+    inline void
+    Release(const uint8_t* data) const {
+        return cast().ReleaseImpl(data);
+    }
+
 private:
     inline IOTmpl&
     cast() {
diff --git a/src/io/io_read_object.h b/src/io/io_read_object.h
new file mode 100644
index 00000000..1879b686
--- /dev/null
+++ b/src/io/io_read_object.h
@@ -0,0 +1,37 @@
+
+// Copyright 2024-present the vsag project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <cstdint>
+
+#include "basic_io.h"
+
+class IOReadObject {
+public:
+    explicit IOReadObject(std::function<void(const uint8_t*)> func) : release_func_(func){};
+
+    const uint8_t* data_;
+
+    bool need_release_{false};
+
+    std::function<void(const uint8_t*)> release_func_{};
+
+    ~IOReadObject() {
+        if (need_release_) {
+            release_func_(data_);
+        }
+    }
+};
\ No newline at end of file
diff --git a/src/io/memory_block_io.h b/src/io/memory_block_io.h
index 3e045254..168b7a62 100644
--- a/src/io/memory_block_io.h
+++ b/src/io/memory_block_io.h
@@ -59,10 +59,16 @@ class MemoryBlockIO : public BasicIO<MemoryBlockIO> {
     WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset);
 
     inline bool
-    ReadImpl(uint8_t* data, uint64_t size, uint64_t offset) const;
+    ReadImpl(uint64_t size, uint64_t offset, uint8_t* data) const;
 
     [[nodiscard]] inline const uint8_t*
-    ReadImpl(uint64_t size, uint64_t offset) const;
+    ReadImpl(uint64_t size, uint64_t offset, bool& need_release) const;
+
+    inline void
+    ReleaseImpl(const uint8_t* data) const {
+        auto ptr = const_cast<uint8_t*>(data);
+        allocator_->Deallocate(ptr);
+    };
 
     inline bool
     MultiReadImpl(uint8_t* datas, uint64_t* sizes, uint64_t* offsets, uint64_t count) const;
@@ -126,7 +132,7 @@ MemoryBlockIO::WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset) {
 }
 
 bool
-MemoryBlockIO::ReadImpl(uint8_t* data, uint64_t size, uint64_t offset) const {
+MemoryBlockIO::ReadImpl(uint64_t size, uint64_t offset, uint8_t* data) const {
     bool ret = check_valid_offset(size + offset);
     if (ret) {
         uint64_t cur_size = 0;
@@ -147,14 +153,16 @@ MemoryBlockIO::ReadImpl(uint8_t* data, uint64_t size, uint64_t offset) const {
 }
 
 const uint8_t*
-MemoryBlockIO::ReadImpl(uint64_t size, uint64_t offset) const {
+MemoryBlockIO::ReadImpl(uint64_t size, uint64_t offset, bool& need_release) const {
     if (check_valid_offset(size + offset)) {
         if (check_in_one_block(offset, size + offset)) {
+            need_release = false;
             return this->get_data_ptr(offset);
         } else {
-            throw std::runtime_error(
-                "Memory Block IO don't support read cross-block data without passing data "
-                "pointer");
+            need_release = true;
+            uint8_t* ptr = reinterpret_cast<uint8_t*>(allocator_->Allocate(size));
+            this->ReadImpl(size, offset, ptr);
+            return ptr;
         }
     }
     return nullptr;
@@ -166,7 +174,7 @@ MemoryBlockIO::MultiReadImpl(uint8_t* datas,
                              uint64_t count) const {
     bool ret = true;
     for (uint64_t i = 0; i < count; ++i) {
-        ret &= this->ReadImpl(datas, sizes[i], offsets[i]);
+        ret &= this->ReadImpl(sizes[i], offsets[i], datas);
         datas += sizes[i];
     }
     return ret;
diff --git a/src/io/memory_io.h b/src/io/memory_io.h
index 81d3434f..545afafc 100644
--- a/src/io/memory_io.h
+++ b/src/io/memory_io.h
@@ -49,10 +49,13 @@ class MemoryIO : public BasicIO<MemoryIO> {
     WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset);
 
     inline bool
-    ReadImpl(uint8_t* data, uint64_t size, uint64_t offset) const;
+    ReadImpl(uint64_t size, uint64_t offset, uint8_t* data) const;
 
     [[nodiscard]] inline const uint8_t*
-    ReadImpl(uint64_t size, uint64_t offset) const;
+    ReadImpl(uint64_t size, uint64_t offset, bool& need_release) const;
+
+    inline void
+    ReleaseImpl(const uint8_t* data) const {};
 
     inline bool
     MultiReadImpl(uint8_t* datas, uint64_t* sizes, uint64_t* offsets, uint64_t count) const;
@@ -95,7 +98,7 @@ MemoryIO::WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset) {
 }
 
 bool
-MemoryIO::ReadImpl(uint8_t* data, uint64_t size, uint64_t offset) const {
+MemoryIO::ReadImpl(uint64_t size, uint64_t offset, uint8_t* data) const {
     bool ret = checkValidOffset(size + offset);
     if (ret) {
         memcpy(data, start_ + offset, size);
@@ -104,7 +107,8 @@ MemoryIO::ReadImpl(uint8_t* data, uint64_t size, uint64_t offset) const {
 }
 
 const uint8_t*
-MemoryIO::ReadImpl(uint64_t size, uint64_t offset) const {
+MemoryIO::ReadImpl(uint64_t size, uint64_t offset, bool& need_release) const {
+    need_release = false;
     if (checkValidOffset(size + offset)) {
         return start_ + offset;
     }
@@ -114,7 +118,7 @@ bool
 MemoryIO::MultiReadImpl(uint8_t* datas, uint64_t* sizes, uint64_t* offsets, uint64_t count) const {
     bool ret = true;
     for (uint64_t i = 0; i < count; ++i) {
-        ret &= this->ReadImpl(datas, sizes[i], offsets[i]);
+        ret &= this->ReadImpl(sizes[i], offsets[i], datas);
         datas += sizes[i];
     }
     return ret;


Test patch:
diff --git a/src/io/basic_io_test.h b/src/io/basic_io_test.h
index ec6db016..a4bfeba3 100644
--- a/src/io/basic_io_test.h
+++ b/src/io/basic_io_test.h
@@ -29,7 +29,7 @@ template <typename T>
 void
 TestBasicReadWrite(BasicIO<T>& io) {
     std::vector<uint64_t> counts = {200, 500};
-    std::vector<uint64_t> max_lengths = {2, 20, 37, 64, 128, 260, 999, 4097};
+    std::vector<uint64_t> max_lengths = {2, 20, 37, 64, 128, 260, 999};
     for (auto count : counts) {
         for (auto max_length : max_lengths) {
             auto vecs = fixtures::GenTestItems(count, max_length);
@@ -38,8 +38,14 @@ TestBasicReadWrite(BasicIO<T>& io) {
             }
             for (auto& item : vecs) {
                 std::vector<uint8_t> data(item.length_);
-                io.Read(data.data(), item.length_, item.start_);
+                io.Read(item.length_, item.start_, data.data());
+                bool need_release = false;
+                const auto* ptr = io.Read(item.length_, item.start_, need_release);
+                REQUIRE(memcmp(data.data(), ptr, item.length_) == 0);
                 REQUIRE(memcmp(data.data(), item.data_, item.length_) == 0);
+                if (need_release) {
+                    io.Release(ptr);
+                }
             }
         }
     }
@@ -49,7 +55,7 @@ template <typename T>
 void
 TestSerializeAndDeserialize(BasicIO<T>& wio, BasicIO<T>& rio) {
     std::vector<uint64_t> counts = {200, 500};
-    std::vector<uint64_t> max_lengths = {2, 20, 37, 64, 128, 260, 999, 4097};
+    std::vector<uint64_t> max_lengths = {2, 20, 37, 64, 128, 260, 999};
     srandom(time(nullptr));
     fixtures::temp_dir dirname("TestSerializeAndDeserialize");
     for (auto count : counts) {
@@ -70,8 +76,14 @@ TestSerializeAndDeserialize(BasicIO<T>& wio, BasicIO<T>& rio) {
 
             for (auto& item : vecs) {
                 std::vector<uint8_t> data(item.length_);
-                rio.Read(data.data(), item.length_, item.start_);
+                rio.Read(item.length_, item.start_, data.data());
+                bool need_release = false;
+                const auto* ptr = rio.Read(item.length_, item.start_, need_release);
+                REQUIRE(memcmp(data.data(), ptr, item.length_) == 0);
                 REQUIRE(memcmp(data.data(), item.data_, item.length_) == 0);
+                if (need_release) {
+                    rio.Release(ptr);
+                }
             }
             infile.close();
         }
diff --git a/src/io/memory_block_io_test.cpp b/src/io/memory_block_io_test.cpp
index b5f940ce..1188d0f3 100644
--- a/src/io/memory_block_io_test.cpp
+++ b/src/io/memory_block_io_test.cpp
@@ -22,15 +22,21 @@
 #include "default_allocator.h"
 using namespace vsag;
 
+auto block_memory_io_block_sizes = {64, 1023, 4096, 123123, 1024 * 1024 * 1024};
+
 TEST_CASE("read&write [ut][memory_block_io]") {
     auto allocator = std::make_unique<DefaultAllocator>();
-    auto io = std::make_unique<MemoryBlockIO>(allocator.get());
-    TestBasicReadWrite(*io);
+    for (auto block_size : block_memory_io_block_sizes) {
+        auto io = std::make_unique<MemoryBlockIO>(allocator.get(), block_size);
+        TestBasicReadWrite(*io);
+    }
 }
 
 TEST_CASE("serialize&deserialize [ut][memory_block_io]") {
     auto allocator = std::make_unique<DefaultAllocator>();
-    auto wio = std::make_unique<MemoryBlockIO>(allocator.get());
-    auto rio = std::make_unique<MemoryBlockIO>(allocator.get());
-    TestSerializeAndDeserialize(*wio, *rio);
+    for (auto block_size : block_memory_io_block_sizes) {
+        auto wio = std::make_unique<MemoryBlockIO>(allocator.get(), block_size);
+        auto rio = std::make_unique<MemoryBlockIO>(allocator.get(), block_size);
+        TestSerializeAndDeserialize(*wio, *rio);
+    }
 }