[MOD-9557] Fix incorrect vector blob size calculation

cmake/svs.cmake

@@ -62,8 +62,8 @@ if(USE_SVS)
         find_package(svs REQUIRED)
         set(SVS_LVQ_HEADER "svs/extensions/vamana/lvq.h")
     else()
-        # This file is included from both CMakeLists.txt and python_bindings/CMakeLists.txt  
-        # Set `root` relative to this file, regardless of where it is included from. 
+        # This file is included from both CMakeLists.txt and python_bindings/CMakeLists.txt
+        # Set `root` relative to this file, regardless of where it is included from.
         get_filename_component(root ${CMAKE_CURRENT_LIST_DIR}/.. ABSOLUTE)
         add_subdirectory(
             ${root}/deps/ScalableVectorSearch

src/VecSim/algorithms/brute_force/brute_force.h

@@ -42,8 +42,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;
@@ -76,16 +76,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;
@@ -350,12 +340,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

@@ -47,10 +47,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

@@ -49,9 +49,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

@@ -302,16 +302,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();
@@ -1561,7 +1551,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

@@ -74,10 +74,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;
 
@@ -201,13 +219,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

@@ -50,9 +50,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;
 
@@ -161,13 +176,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,10 @@ 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.
+    : 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 +1192,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/algorithms/svs/svs.h

@@ -155,11 +155,12 @@
             return MemoryUtils::unique_blob{const_cast<void *>(original_data), [](void *) {}};
         }
 
-        const auto data_size = this->dim * sizeof(DataType) * n;
+        const auto data_size = this->getDataSize() * n;
 
         auto processed_blob =
             MemoryUtils::unique_blob{this->allocator->allocate(data_size),
                                      [this](void *ptr) { this->allocator->free_allocation(ptr); }};
+        // Assuming original data size equals to processed data size
         memcpy(processed_blob.get(), original_data, data_size);
         // Preprocess each vector in place
         for (size_t i = 0; i < n; i++) {
@@ -435,17 +436,18 @@
 
     VecSimBatchIterator *newBatchIterator(const void *queryBlob,
                                           VecSimQueryParams *queryParams) const override {
-        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 VecSimBatchIterator that will free it at the end.
         if (indexSize() == 0) {
             return new (this->getAllocator())
-                NullSVS_BatchIterator(queryBlobCopy, queryParams, this->getAllocator());
+                NullSVS_BatchIterator(queryBlobCopyPtr, queryParams, this->getAllocator());
         } else {
             return new (this->getAllocator()) SVS_BatchIterator<impl_type, data_type>(
-                queryBlobCopy, impl_.get(), queryParams, this->getAllocator());
+                queryBlobCopyPtr, impl_.get(), queryParams, this->getAllocator());
         }
     }
 
@@ -479,6 +481,15 @@
     }
 
 #ifdef BUILD_TESTS
-    virtual void fitMemory() {};
+    void fitMemory() override {}
+    std::vector<std::vector<char>> getStoredVectorDataByLabel(labelType label) const override {
+        assert(nullptr && "Not implemented");
+        return {};
+    }
+    void getDataByLabel(
+        labelType label,
+        std::vector<std::vector<svs_details::vecsim_dt<DataType>>> &vectors_output) const override {
+        assert(nullptr && "Not implemented");
+    }
 #endif
 };