[8.0] [MOD-9557] Fix incorrect vector blob size calculation

src/VecSim/algorithms/brute_force/brute_force.h

@@ -43,8 +43,8 @@ class BruteForceIndex : public VecSimIndexAbstract<DataType, DistType> {
     size_t indexSize() const override;
     size_t indexCapacity() const override;
     std::unique_ptr<RawDataContainer::Iterator> getVectorsIterator() const;
-    DataType *getDataByInternalId(idType id) const {
-        return (DataType *)this->vectors->getElement(id);
+    const DataType *getDataByInternalId(idType id) const {
+        return reinterpret_cast<const DataType *>(this->vectors->getElement(id));
     }
     VecSimQueryReply *topKQuery(const void *queryBlob, size_t k,
                                 VecSimQueryParams *queryParams) const override;
@@ -77,16 +77,6 @@ class BruteForceIndex : public VecSimIndexAbstract<DataType, DistType> {
 
     virtual ~BruteForceIndex() = default;
 #ifdef BUILD_TESTS
-    /**
-     * @brief Used for testing - store vector(s) data associated with a given label. This function
-     * copies the vector(s)' data buffer(s) and place it in the output vector
-     *
-     * @param label
-     * @param vectors_output empty vector to be modified, should store the blob(s) associated with
-     * the label.
-     */
-    virtual void getDataByLabel(labelType label,
-                                std::vector<std::vector<DataType>> &vectors_output) const = 0;
     void fitMemory() override {
         if (count == 0) {
             return;
@@ -351,12 +341,13 @@ template <typename DataType, typename DistType>
 VecSimBatchIterator *
 BruteForceIndex<DataType, DistType>::newBatchIterator(const void *queryBlob,
                                                       VecSimQueryParams *queryParams) const {
-    auto *queryBlobCopy =
-        this->allocator->allocate_aligned(this->dataSize, this->preprocessors->getAlignment());
-    memcpy(queryBlobCopy, queryBlob, this->dim * sizeof(DataType));
-    this->preprocessQueryInPlace(queryBlobCopy);
+    // force_copy == true.
+    auto queryBlobCopy = this->preprocessQuery(queryBlob, true);
+
+    // take ownership of the blob copy and pass it to the batch iterator.
+    auto *queryBlobCopyPtr = queryBlobCopy.release();
     // Ownership of queryBlobCopy moves to BF_BatchIterator that will free it at the end.
-    return newBatchIterator_Instance(queryBlobCopy, queryParams);
+    return newBatchIterator_Instance(queryBlobCopyPtr, queryParams);
 }
 
 template <typename DataType, typename DistType>

src/VecSim/algorithms/brute_force/brute_force_multi.h

@@ -48,10 +48,30 @@
 
         for (idType id : ids->second) {
             auto vec = std::vector<DataType>(this->dim);
+            // Only copy the vector data (dim * sizeof(DataType)), not any additional metadata like
+            // the norm
             memcpy(vec.data(), this->getDataByInternalId(id), this->dim * sizeof(DataType));
             vectors_output.push_back(vec);
         }
     }
+
+    std::vector<std::vector<char>> getStoredVectorDataByLabel(labelType label) const override {
+        std::vector<std::vector<char>> vectors_output;
+        auto ids = labelToIdsLookup.find(label);
+
+        for (idType id : ids->second) {
+            // Get the data pointer - need to cast to char* for memcpy
+            const char *data = reinterpret_cast<const char *>(this->getDataByInternalId(id));
+
+            // Create a vector with the full data (including any metadata like norms)
+            std::vector<char> vec(this->getDataSize());
+            memcpy(vec.data(), data, this->getDataSize());
+            vectors_output.push_back(std::move(vec));
+        }
+
+        return vectors_output;
+    }
+
 #endif
 private:
     // inline definitions

src/VecSim/algorithms/brute_force/brute_force_single.h

@@ -50,9 +50,26 @@ class BruteForceIndex_Single : public BruteForceIndex<DataType, DistType> {
         auto id = labelToIdLookup.at(label);
 
         auto vec = std::vector<DataType>(this->dim);
+        // Only copy the vector data (dim * sizeof(DataType)), not any additional metadata like the
+        // norm
         memcpy(vec.data(), this->getDataByInternalId(id), this->dim * sizeof(DataType));
         vectors_output.push_back(vec);
     }
+
+    std::vector<std::vector<char>> getStoredVectorDataByLabel(labelType label) const override {
+        std::vector<std::vector<char>> vectors_output;
+        auto id = labelToIdLookup.at(label);
+
+        // Get the data pointer - need to cast to char* for memcpy
+        const char *data = reinterpret_cast<const char *>(this->getDataByInternalId(id));
+
+        // Create a vector with the full data (including any metadata like norms)
+        std::vector<char> vec(this->getDataSize());
+        memcpy(vec.data(), data, this->getDataSize());
+        vectors_output.push_back(std::move(vec));
+
+        return vectors_output;
+    }
 #endif
 protected:
     // inline definitions

src/VecSim/algorithms/hnsw/hnsw.h

@@ -303,16 +303,6 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
     virtual int removeLabel(labelType label) = 0;
 
 #ifdef BUILD_TESTS
-    /**
-     * @brief Used for testing - store vector(s) data associated with a given label. This function
-     * copies the vector(s)' data buffer(s) and place it in the output vector
-     *
-     * @param label
-     * @param vectors_output empty vector to be modified, should store the blob(s) associated with
-     * the label.
-     */
-    virtual void getDataByLabel(labelType label,
-                                std::vector<std::vector<DataType>> &vectors_output) const = 0;
     void fitMemory() override {
         if (maxElements > 0) {
             idToMetaData.shrink_to_fit();
@@ -1562,7 +1552,7 @@ void HNSWIndex<DataType, DistType>::insertElementToGraph(idType element_id,
     for (auto level = static_cast<int>(max_common_level); level >= 0; level--) {
         candidatesMaxHeap<DistType> top_candidates =
             searchLayer(curr_element, vector_data, level, efConstruction);
-        // If the entry point was marked deleted between iterations, we may recieve an empty
+        // If the entry point was marked deleted between iterations, we may receive an empty
         // candidates set.
         if (!top_candidates.empty()) {
             curr_element = mutuallyConnectNewElement(element_id, top_candidates, level);

src/VecSim/algorithms/hnsw/hnsw_multi.h

@@ -75,10 +75,28 @@
 
         for (idType id : ids->second) {
             auto vec = std::vector<DataType>(this->dim);
-            memcpy(vec.data(), this->getDataByInternalId(id), this->dataSize);
+            // Only copy the vector data (dim * sizeof(DataType)), not any additional metadata like
+            // the norm
+            memcpy(vec.data(), this->getDataByInternalId(id), this->dim * sizeof(DataType));
             vectors_output.push_back(vec);
         }
     }
+
+    std::vector<std::vector<char>> getStoredVectorDataByLabel(labelType label) const override {
+        std::vector<std::vector<char>> vectors_output;
+        auto ids = labelLookup.find(label);
+
+        for (idType id : ids->second) {
+            const char *data = this->getDataByInternalId(id);
+
+            // Create a vector with the full data (including any metadata like norms)
+            std::vector<char> vec(this->dataSize);
+            memcpy(vec.data(), data, this->dataSize);
+            vectors_output.push_back(std::move(vec));
+        }
+
+        return vectors_output;
+    }
 #endif
     ~HNSWIndex_Multi() = default;
 
@@ -202,13 +220,14 @@
 VecSimBatchIterator *
 HNSWIndex_Multi<DataType, DistType>::newBatchIterator(const void *queryBlob,
                                                       VecSimQueryParams *queryParams) const {
-    auto queryBlobCopy =
-        this->allocator->allocate_aligned(this->dataSize, this->preprocessors->getAlignment());
-    memcpy(queryBlobCopy, queryBlob, this->dim * sizeof(DataType));
-    this->preprocessQueryInPlace(queryBlobCopy);
+    // force_copy == true.
+    auto queryBlobCopy = this->preprocessQuery(queryBlob, true);
+
+    // take ownership of the blob copy and pass it to the batch iterator.
+    auto *queryBlobCopyPtr = queryBlobCopy.release();
     // Ownership of queryBlobCopy moves to HNSW_BatchIterator that will free it at the end.
     return new (this->allocator) HNSWMulti_BatchIterator<DataType, DistType>(
-        queryBlobCopy, this, queryParams, this->allocator);
+        queryBlobCopyPtr, this, queryParams, this->allocator);
 }
 
 /**

src/VecSim/algorithms/hnsw/hnsw_single.h

@@ -51,9 +51,24 @@ class HNSWIndex_Single : public HNSWIndex<DataType, DistType> {
         auto id = labelLookup.at(label);
 
         auto vec = std::vector<DataType>(this->dim);
-        memcpy(vec.data(), this->getDataByInternalId(id), this->dataSize);
+        // Only copy the vector data (dim * sizeof(DataType)), not any additional metadata like the
+        // norm
+        memcpy(vec.data(), this->getDataByInternalId(id), this->dim * sizeof(DataType));
         vectors_output.push_back(vec);
     }
+
+    std::vector<std::vector<char>> getStoredVectorDataByLabel(labelType label) const override {
+        std::vector<std::vector<char>> vectors_output;
+        auto id = labelLookup.at(label);
+        const char *data = this->getDataByInternalId(id);
+
+        // Create a vector with the full data (including any metadata like norms)
+        std::vector<char> vec(this->dataSize);
+        memcpy(vec.data(), data, this->dataSize);
+        vectors_output.push_back(std::move(vec));
+
+        return vectors_output;
+    }
 #endif
     ~HNSWIndex_Single() = default;
 
@@ -162,13 +177,14 @@ template <typename DataType, typename DistType>
 VecSimBatchIterator *
 HNSWIndex_Single<DataType, DistType>::newBatchIterator(const void *queryBlob,
                                                        VecSimQueryParams *queryParams) const {
-    auto queryBlobCopy =
-        this->allocator->allocate_aligned(this->dataSize, this->preprocessors->getAlignment());
-    memcpy(queryBlobCopy, queryBlob, this->dim * sizeof(DataType));
-    this->preprocessQueryInPlace(queryBlobCopy);
+    // force_copy == true.
+    auto queryBlobCopy = this->preprocessQuery(queryBlob, true);
+
+    // take ownership of the blob copy and pass it to the batch iterator.
+    auto *queryBlobCopyPtr = queryBlobCopy.release();
     // Ownership of queryBlobCopy moves to HNSW_BatchIterator that will free it at the end.
     return new (this->allocator) HNSWSingle_BatchIterator<DataType, DistType>(
-        queryBlobCopy, this, queryParams, this->allocator);
+        queryBlobCopyPtr, this, queryParams, this->allocator);
 }
 
 /**

src/VecSim/algorithms/hnsw/hnsw_tiered.h

@@ -172,7 +172,7 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
         inline void filter_irrelevant_results(VecSimQueryResultContainer &);
 
     public:
-        TieredHNSW_BatchIterator(void *query_vector,
+        TieredHNSW_BatchIterator(const void *query_vector,
                                  const TieredHNSWIndex<DataType, DistType> *index,
                                  VecSimQueryParams *queryParams,
                                  std::shared_ptr<VecSimAllocator> allocator);
@@ -206,11 +206,9 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
     VecSimDebugInfoIterator *debugInfoIterator() const override;
     VecSimBatchIterator *newBatchIterator(const void *queryBlob,
                                           VecSimQueryParams *queryParams) const override {
-        size_t blobSize = this->frontendIndex->getDim() * sizeof(DataType);
-        void *queryBlobCopy = this->allocator->allocate(blobSize);
-        memcpy(queryBlobCopy, queryBlob, blobSize);
+        // The query blob will be processed and copied by the internal indexes's batch iterator.
         return new (this->allocator)
-            TieredHNSW_BatchIterator(queryBlobCopy, this, queryParams, this->allocator);
+            TieredHNSW_BatchIterator(queryBlob, this, queryParams, this->allocator);
     }
     inline void setLastSearchMode(VecSearchMode mode) override {
         return this->backendIndex->setLastSearchMode(mode);
@@ -545,10 +543,11 @@ void TieredHNSWIndex<DataType, DistType>::executeInsertJob(HNSWInsertJob *job) {
     HNSWIndex<DataType, DistType> *hnsw_index = this->getHNSWIndex();
     // Copy the vector blob from the flat buffer, so we can release the flat lock while we are
     // indexing the vector into HNSW index.
-    auto blob_copy = this->getAllocator()->allocate_unique(this->frontendIndex->getDataSize());
-
-    memcpy(blob_copy.get(), this->frontendIndex->getDataByInternalId(job->id),
-           this->frontendIndex->getDim() * sizeof(DataType));
+    size_t data_size = this->frontendIndex->getDataSize();
+    auto blob_copy = this->getAllocator()->allocate_unique(data_size);
+    // Assuming the size of the blob stored in the frontend index matches the size of the blob
+    // stored in the HNSW index.
+    memcpy(blob_copy.get(), this->frontendIndex->getDataByInternalId(job->id), data_size);
 
     this->insertVectorToHNSW<true>(hnsw_index, job->label, blob_copy.get());
 
@@ -719,7 +718,7 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
     int ret = 1;
     auto hnsw_index = this->getHNSWIndex();
     // writeMode is not protected since it is assumed to be called only from the "main thread"
-    // (that is the thread that is exculusively calling add/delete vector).
+    // (that is the thread that is exclusively calling add/delete vector).
     if (this->getWriteMode() == VecSim_WriteInPlace) {
         // First, check if we need to overwrite the vector in-place for single (from both indexes).
         if (!this->backendIndex->isMultiValue()) {
@@ -849,7 +848,7 @@ int TieredHNSWIndex<DataType, DistType>::deleteVector(labelType label) {
     // Note that we may remove the same vector that has been removed from the flat index, if it was
     // being ingested at that time.
     // writeMode is not protected since it is assumed to be called only from the "main thread"
-    // (that is the thread that is exculusively calling add/delete vector).
+    // (that is the thread that is exclusively calling add/delete vector).
     if (this->getWriteMode() == VecSim_WriteAsync) {
         num_deleted_vectors += this->deleteLabelFromHNSW(label);
         // Apply ready swap jobs if number of deleted vectors reached the threshold
@@ -924,9 +923,14 @@ double TieredHNSWIndex<DataType, DistType>::getDistanceFrom_Unsafe(labelType lab
 
 template <typename DataType, typename DistType>
 TieredHNSWIndex<DataType, DistType>::TieredHNSW_BatchIterator::TieredHNSW_BatchIterator(
-    void *query_vector, const TieredHNSWIndex<DataType, DistType> *index,
+    const void *query_vector, const TieredHNSWIndex<DataType, DistType> *index,
     VecSimQueryParams *queryParams, std::shared_ptr<VecSimAllocator> allocator)
-    : VecSimBatchIterator(query_vector, queryParams ? queryParams->timeoutCtx : nullptr,
+    // Tiered batch iterator doesn't hold its own copy of the query vector.
+    // Instead, each internal batch iterators (flat_iterator and hnsw_iterator) create their own
+    // copies: flat_iterator copy is created during TieredHNSW_BatchIterator construction When
+    // TieredHNSW_BatchIterator::getNextResults() is called and hnsw_iterator is not initialized, it
+    // retrieves the blob from flat_iterator
+    : VecSimBatchIterator(nullptr, queryParams ? queryParams->timeoutCtx : nullptr,
                           std::move(allocator)),
       index(index), flat_results(this->allocator), hnsw_results(this->allocator),
       flat_iterator(this->index->frontendIndex->newBatchIterator(query_vector, queryParams)),
@@ -1192,4 +1196,5 @@ void TieredHNSWIndex<DataType, DistType>::getDataByLabel(
     labelType label, std::vector<std::vector<DataType>> &vectors_output) const {
     this->getHNSWIndex()->getDataByLabel(label, vectors_output);
 }
+
 #endif

src/VecSim/spaces/computer/preprocessor_container.cpp

@@ -21,10 +21,9 @@ PreprocessorsContainerAbstract::preprocessForStorage(const void *original_blob,
     return wrapWithDummyDeleter(const_cast<void *>(original_blob));
 }
 
-MemoryUtils::unique_blob
-PreprocessorsContainerAbstract::preprocessQuery(const void *original_blob,
-                                                size_t processed_bytes_count) const {
-    return maybeCopyToAlignedMem(original_blob, processed_bytes_count);
+MemoryUtils::unique_blob PreprocessorsContainerAbstract::preprocessQuery(
+    const void *original_blob, size_t processed_bytes_count, bool force_copy) const {
+    return maybeCopyToAlignedMem(original_blob, processed_bytes_count, force_copy);
 }
 
 void PreprocessorsContainerAbstract::preprocessQueryInPlace(void *blob,
@@ -33,15 +32,16 @@ void PreprocessorsContainerAbstract::preprocessQueryInPlace(void *blob,
 void PreprocessorsContainerAbstract::preprocessStorageInPlace(void *blob,
                                                               size_t processed_bytes_count) const {}
 
-MemoryUtils::unique_blob
-PreprocessorsContainerAbstract::maybeCopyToAlignedMem(const void *original_blob,
-                                                      size_t blob_bytes_count) const {
-    if (this->alignment) {
-        if ((uintptr_t)original_blob % this->alignment) {
-            auto aligned_mem = this->allocator->allocate_aligned(blob_bytes_count, this->alignment);
-            memcpy(aligned_mem, original_blob, blob_bytes_count);
-            return this->wrapAllocated(aligned_mem);
-        }
+MemoryUtils::unique_blob PreprocessorsContainerAbstract::maybeCopyToAlignedMem(
+    const void *original_blob, size_t blob_bytes_count, bool force_copy) const {
+    bool needs_copy =
+        force_copy || (this->alignment && ((uintptr_t)original_blob % this->alignment != 0));
+
+    if (needs_copy) {
+        auto aligned_mem = this->allocator->allocate_aligned(blob_bytes_count, this->alignment);
+        // TODO: handle original_blob_size != processed_bytes_count
+        memcpy(aligned_mem, original_blob, blob_bytes_count);
+        return this->wrapAllocated(aligned_mem);
     }
 
     // Returning a unique_ptr with a no-op deleter