diff --git a/AnnService/CMakeLists.txt b/AnnService/CMakeLists.txt index ecc01eb12..b48354978 100644 --- a/AnnService/CMakeLists.txt +++ b/AnnService/CMakeLists.txt @@ -144,17 +144,23 @@ if (RABITQ) else() list(REMOVE_ITEM HDR_FILES ${AnnService}/inc/Core/Common/RaBitQQuantizer.h + ${AnnService}/inc/Core/Common/RaBitQAutoTuner.h ) list(REMOVE_ITEM SRC_FILES ${AnnService}/src/Core/Common/RaBitQQuantizer.cpp + ${AnnService}/src/Core/Common/RaBitQAutoTuner.cpp ) endif() add_library (SPTAGLib SHARED ${SRC_FILES} ${HDR_FILES} ${TiKV_PROTO_SOURCES}) -target_link_libraries (SPTAGLib DistanceUtils ${RabitQ_LIBRARIES} ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_shared ${NUMA_LIBRARY} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES}) +target_link_libraries (SPTAGLib DistanceUtils ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_shared ${NUMA_LIBRARY} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES}) add_library (SPTAGLibStatic STATIC ${SRC_FILES} ${HDR_FILES} ${TiKV_PROTO_SOURCES}) -target_link_libraries (SPTAGLibStatic DistanceUtils ${RabitQ_LIBRARIES} ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_static ${NUMA_LIBRARY_STATIC} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES}) +target_link_libraries (SPTAGLibStatic DistanceUtils ${RocksDB_LIBRARIES} ${uring_LIBRARIES} libzstd_static ${NUMA_LIBRARY_STATIC} ${TBB_LIBRARIES} ${SPDK_LIBRARIES} ${TiKV_LIBRARIES}) +if (RABITQ) + target_link_libraries(SPTAGLib RaBitQOfficialCore) + target_link_libraries(SPTAGLibStatic RaBitQOfficialCore) +endif() if (MSVC) # SPANNIndex.cpp can exceed COFF section limits in Debug without /bigobj. diff --git a/AnnService/inc/Core/Common/IQuantizer.h b/AnnService/inc/Core/Common/IQuantizer.h index 1dde003c4..7c203b80f 100644 --- a/AnnService/inc/Core/Common/IQuantizer.h +++ b/AnnService/inc/Core/Common/IQuantizer.h @@ -60,6 +60,8 @@ namespace SPTAG virtual float* GetL2DistanceTables() = 0; + virtual bool QuantizeForIndexBuild() const { return true; } + template T* GetCodebooks(); }; diff --git a/AnnService/inc/Core/Common/OPQQuantizer.h b/AnnService/inc/Core/Common/OPQQuantizer.h index 9a1f73c55..b839f42ff 100644 --- a/AnnService/inc/Core/Common/OPQQuantizer.h +++ b/AnnService/inc/Core/Common/OPQQuantizer.h @@ -53,6 +53,10 @@ namespace SPTAG return GetEnumValueType(); } + bool QuantizeForIndexBuild() const override + { + return false; + } protected: using PQQuantizer::m_NumSubvectors; @@ -200,7 +204,11 @@ namespace SPTAG inline void OPQQuantizer::m_VectorMatrixMultiply(OPQMatrixType* mat, const OPQMatrixType* vec, O* mat_vec) const { for (int i = 0; i < m_matrixDim; i++) { - mat_vec[i] = (O)(m_base - m_fdot(vec, mat, m_matrixDim)); + OPQMatrixType value = 0; + for (int j = 0; j < m_matrixDim; ++j) { + value += vec[j] * mat[j]; + } + mat_vec[i] = static_cast(value); mat += m_matrixDim; } } diff --git a/AnnService/inc/Core/Common/RaBitQAutoTuner.h b/AnnService/inc/Core/Common/RaBitQAutoTuner.h new file mode 100644 index 000000000..9bddac8dc --- /dev/null +++ b/AnnService/inc/Core/Common/RaBitQAutoTuner.h @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "inc/Core/Common/RaBitQQuantizer.h" +#include "inc/Helper/SimpleIniReader.h" +#include "inc/Helper/VectorSetReader.h" + +#include +#include +#include + +namespace SPTAG +{ +namespace COMMON +{ + +struct RaBitQAutoTuneResult +{ + int selectedBits = 0; + float recall = 0.0F; + SizeType vectorCount = 0; + DimensionType codeDimension = 0; + std::string quantizerPath; + std::string vectorPath; + std::shared_ptr quantizer; +}; + +class RaBitQAutoTuner +{ +public: + using BitEvaluator = std::function; + + static bool IsEnabled(const Helper::IniReader& p_config); + static ErrorCode Run(Helper::IniReader& p_config, + const std::string& p_outputFolder, + RaBitQAutoTuneResult& p_result, + std::string& p_error); + + static ErrorCode SelectMinimumBits(float p_targetRecall, + const BitEvaluator& p_evaluator, + int& p_selectedBits, + float& p_selectedRecall); + static ErrorCode ValidateTruth(const std::vector>& p_truth, + SizeType p_baseCount, + SizeType p_queryCount, + int p_resultCount, + std::string& p_error); + static float RecallAtK(const std::vector& p_exact, + const std::vector& p_reranked, + int p_resultCount); +}; + +} // namespace COMMON +} // namespace SPTAG diff --git a/AnnService/inc/Core/Common/RaBitQQuantizer.h b/AnnService/inc/Core/Common/RaBitQQuantizer.h index dd036be31..612f9f626 100644 --- a/AnnService/inc/Core/Common/RaBitQQuantizer.h +++ b/AnnService/inc/Core/Common/RaBitQQuantizer.h @@ -25,6 +25,10 @@ class RaBitQQuantizer : public IQuantizer RaBitQQuantizer(DimensionType p_dimension, int p_bits, bool p_normalize); ErrorCode Train(const std::shared_ptr& p_vectors); + ErrorCode BeginTraining(); + ErrorCode AddTrainingBatch(const std::shared_ptr& p_vectors); + ErrorCode FinishTraining(); + std::shared_ptr CreateWithBits(int p_bits) const; float L2Distance(const std::uint8_t* p_x, const std::uint8_t* p_y) const override; float CosineDistance(const std::uint8_t* p_x, const std::uint8_t* p_y) const override; @@ -44,10 +48,12 @@ class RaBitQQuantizer : public IQuantizer DimensionType GetNumSubvectors() const override; int GetBase() const override; float* GetL2DistanceTables() override; + bool QuantizeForIndexBuild() const override { return false; } DimensionType Dimension() const { return m_dimension; } int Bits() const { return m_bits; } bool Ready() const; + bool Trained() const { return m_trained; } private: struct ModelHeader @@ -90,6 +96,9 @@ class RaBitQQuantizer : public IQuantizer rabitqlib::quant::RabitqConfig m_quantizer_config; rabitqlib::ex_ipfunc m_ip_func = nullptr; std::vector m_centroid; + std::vector m_training_sum; + std::uint64_t m_training_count = 0; + bool m_trained = false; }; } // namespace COMMON diff --git a/AnnService/inc/Core/Common/TruthSet.h b/AnnService/inc/Core/Common/TruthSet.h index 859e541f4..27f2ab73e 100644 --- a/AnnService/inc/Core/Common/TruthSet.h +++ b/AnnService/inc/Core/Common/TruthSet.h @@ -169,6 +169,11 @@ namespace SPTAG float meanrecall = 0, minrecall = MaxDist, maxrecall = 0, stdrecall = 0, meanmrr = 0; std::vector thisrecall(NumQuerys, 0); std::unique_ptr visited(new bool[K]); + const bool compareDistanceTies = + querySet != nullptr && + vectorSet != nullptr && + querySet->GetValueType() == GetEnumValueType() && + vectorSet->GetValueType() == GetEnumValueType(); for (SizeType i = 0; i < NumQuerys; i++) { int minpos = K; @@ -186,7 +191,7 @@ namespace SPTAG if (j < minpos) minpos = j; break; } - else if (vectorSet != nullptr) { + else if (compareDistanceTies) { float dist = COMMON::DistanceUtils::ComputeDistance((const T*)querySet->GetVector(i), (const T*)vectorSet->GetVector(results[i].GetResult(j)->VID), vectorSet->Dimension(), index->GetDistCalcMethod()); float truthDist = COMMON::DistanceUtils::ComputeDistance((const T*)querySet->GetVector(i), (const T*)vectorSet->GetVector(id), vectorSet->Dimension(), index->GetDistCalcMethod()); if (index->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine && fabs(dist - truthDist) < Epsilon) { @@ -213,7 +218,7 @@ namespace SPTAG std::vector truthvec; for (SizeType id : truth[i]) { float truthDist = 0.0; - if (vectorSet != nullptr) { + if (compareDistanceTies) { truthDist = COMMON::DistanceUtils::ComputeDistance((const T*)querySet->GetVector(i), (const T*)vectorSet->GetVector(id), querySet->Dimension(), index->GetDistCalcMethod()); } truthvec.emplace_back(id, truthDist); diff --git a/AnnService/inc/Core/Common/WorkSpace.h b/AnnService/inc/Core/Common/WorkSpace.h index 6e3404144..59eeb0426 100644 --- a/AnnService/inc/Core/Common/WorkSpace.h +++ b/AnnService/inc/Core/Common/WorkSpace.h @@ -146,13 +146,21 @@ namespace SPTAG inline void DoubleSize() { + const std::uint64_t oldPoolSize = m_poolSize; + const bool hadSecondHash = m_secondHash; std::uint64_t new_poolSize = ((m_poolSize + 1) << 1) - 1; SizeType* new_hashTable = new SizeType[(new_poolSize + 1) * 2]; memset(new_hashTable, 0, sizeof(SizeType) * (new_poolSize + 1) * 2); m_secondHash = false; - for (std::uint64_t i = 0; i <= new_poolSize; i++) + for (std::uint64_t i = 0; i <= oldPoolSize; i++) if (m_hashTable[i]) _CheckAndSet(new_hashTable, new_poolSize, true, m_hashTable[i]); + if (hadSecondHash) + { + SizeType* secondHashTable = m_hashTable.get() + oldPoolSize + 1; + for (std::uint64_t i = 0; i <= oldPoolSize; i++) + if (secondHashTable[i]) _CheckAndSet(new_hashTable, new_poolSize, true, secondHashTable[i]); + } m_exp++; m_poolSize = new_poolSize; diff --git a/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h b/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h index ee01c01f7..112b539da 100644 --- a/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h +++ b/AnnService/inc/Core/SPANN/ExtraStaticSearcher.h @@ -125,8 +125,8 @@ namespace SPTAG SizeType vectorID = *(reinterpret_cast(p_postingListFullData + offsetVectorID));\ if (p_exWorkSpace->Deduper().CheckAndSet(vectorID)) { listElements--; continue; } \ (this->*m_parseEncoding)(listInfo, (ValueType*)(p_postingListFullData + offsetVector));\ - auto distance2leaf = m_headIndex->ComputeDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); \ - queryResults.AddPoint(vectorID, distance2leaf, queryResults.WithVec()? ByteArray((std::uint8_t*)(p_postingListFullData + offsetVector), sizeof(ValueType) * m_opt->m_dim, false) : ByteArray::c_empty); \ + auto distance2leaf = PostingDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); \ + queryResults.AddPoint(vectorID, distance2leaf, queryResults.WithVec()? ByteArray((std::uint8_t*)(p_postingListFullData + offsetVector), StoredVectorBytes(), false) : ByteArray::c_empty); \ } \ #define ProcessPostingOffset() \ @@ -138,7 +138,7 @@ namespace SPTAG if (p_exWorkSpace->Deduper().CheckAndSet(vectorID)) continue; \ if (p_exWorkSpace->m_filterFunc != nullptr && !p_exWorkSpace->m_filterFunc(m_headIndex->GetMetadata(vectorID))) continue; \ (this->*m_parseEncoding)(listInfo, (ValueType*)(p_postingListFullData + offsetVector));\ - auto distance2leaf = m_headIndex->ComputeDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); \ + auto distance2leaf = PostingDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); \ queryResults.AddPoint(vectorID, distance2leaf); \ foundResult = true;\ break;\ @@ -176,6 +176,42 @@ namespace SPTAG return m_available; } + inline size_t StoredVectorBytes() const + { + return m_headIndex->m_pQuantizer + ? static_cast(m_headIndex->m_pQuantizer->GetNumSubvectors()) + : static_cast(m_opt->m_dim) * sizeof(ValueType); + } + + inline float PostingDistance(const void* p_query, const char* p_vector) const + { + if (!m_headIndex->m_pQuantizer) + { + return m_headIndex->ComputeDistance(p_query, p_vector); + } + if (m_headIndex->GetDistCalcMethod() == DistCalcMethod::L2) + { + return m_headIndex->m_pQuantizer->L2Distance( + reinterpret_cast(p_query), + reinterpret_cast(p_vector)); + } + return m_headIndex->m_pQuantizer->CosineDistance( + reinterpret_cast(p_query), + reinterpret_cast(p_vector)); + } + + inline size_t DiskRequestIndex(const ExtraWorkSpace* p_exWorkSpace, size_t p_postingIndex) const + { + return p_postingIndex; + } + + inline Helper::AsyncReadRequest& DiskRequest(ExtraWorkSpace* p_exWorkSpace, size_t p_postingIndex) const + { + auto& request = p_exWorkSpace->m_diskRequests[DiskRequestIndex(p_exWorkSpace, p_postingIndex)]; + request.m_buffer = reinterpret_cast(p_exWorkSpace->m_pageBuffers[p_postingIndex].GetBuffer()); + return request; + } + virtual bool LoadIndex(Options& p_opt) override { m_opt = &p_opt; m_enableDeltaEncoding = p_opt.m_enableDeltaEncoding; @@ -185,7 +221,7 @@ namespace SPTAG m_extraFullGraphFile = p_opt.m_indexDirectory + FolderSep + p_opt.m_ssdIndex; std::string curFile = m_extraFullGraphFile + "_" + std::to_string(m_layer); - p_opt.m_searchPostingPageLimit = max(p_opt.m_searchPostingPageLimit, static_cast((p_opt.m_postingVectorLimit * (p_opt.m_dim * sizeof(ValueType) + sizeof(SizeType)) + PageSize - 1) / PageSize)); + p_opt.m_searchPostingPageLimit = max(p_opt.m_searchPostingPageLimit, static_cast((p_opt.m_postingVectorLimit * (StoredVectorBytes() + sizeof(SizeType)) + PageSize - 1) / PageSize)); SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Load index with posting page limit:%d\n", p_opt.m_searchPostingPageLimit); do { auto curIndexFile = f_createAsyncIO(); @@ -243,7 +279,8 @@ namespace SPTAG { const uint32_t postingListCount = static_cast(p_exWorkSpace->m_postingIDs.size()); if (postingListCount > p_exWorkSpace->m_pageBuffers.size() || - postingListCount > p_exWorkSpace->m_diskRequests.size()) { + (postingListCount > 0 && + DiskRequestIndex(p_exWorkSpace, postingListCount - 1) >= p_exWorkSpace->m_diskRequests.size())) { SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Static search workspace is too small: postings=%u buffers=%zu requests=%zu.\n", postingListCount, p_exWorkSpace->m_pageBuffers.size(), @@ -256,6 +293,7 @@ namespace SPTAG int diskIO = 0; int listElements = 0; int missingPostingIDs = 0; + ErrorCode scanRet = ErrorCode::Success; #if defined(ASYNC_READ) && !defined(BATCH_READ) int unprocessed = 0; @@ -267,7 +305,7 @@ namespace SPTAG auto it = m_globalVectorIDToHeadMap.find(curPostingID); if (it == m_globalVectorIDToHeadMap.end()) { ++missingPostingIDs; - auto& request = p_exWorkSpace->m_diskRequests[pi]; + auto& request = DiskRequest(p_exWorkSpace, pi); request.m_readSize = 0; request.m_success = false; request.m_callback = nullptr; @@ -295,7 +333,7 @@ namespace SPTAG } #ifdef ASYNC_READ - auto& request = p_exWorkSpace->m_diskRequests[pi]; + auto& request = DiskRequest(p_exWorkSpace, pi); request.m_offset = listInfo->listOffset; request.m_readSize = totalBytes; request.m_status = (fileid << 16) | (request.m_status & 0xffff); @@ -303,10 +341,15 @@ namespace SPTAG request.m_success = false; #ifdef BATCH_READ // async batch read - request.m_callback = [&p_exWorkSpace, &queryResults, &request, &listElements, this](bool success) + Helper::AsyncReadRequest* requestPtr = &request; + request.m_callback = [p_exWorkSpace, &queryResults, requestPtr, &listElements, &scanRet, this](bool success) { - char* buffer = request.m_buffer; - ListInfo* listInfo = (ListInfo*)(request.m_payload); + if (!success) { + scanRet = ErrorCode::DiskIOFail; + return; + } + char* buffer = requestPtr->m_buffer; + ListInfo* listInfo = (ListInfo*)(requestPtr->m_payload); // decompress posting list char* p_postingListFullData = buffer + listInfo->pageOffset; @@ -318,9 +361,10 @@ namespace SPTAG ProcessPosting(); }; #else // async read - request.m_callback = [&p_exWorkSpace, &request](bool success) + Helper::AsyncReadRequest* requestPtr = &request; + request.m_callback = [p_exWorkSpace, requestPtr](bool success) { - p_exWorkSpace->m_processIocp.push(&request); + p_exWorkSpace->m_processIocp.push(requestPtr); }; ++unprocessed; @@ -356,7 +400,11 @@ namespace SPTAG #ifdef ASYNC_READ #ifdef BATCH_READ - BatchReadFileAsync(m_indexFiles, (p_exWorkSpace->m_diskRequests).data(), postingListCount); + if (!BatchReadFileAsync(m_indexFiles, (p_exWorkSpace->m_diskRequests).data(), postingListCount)) + { + return ErrorCode::DiskIOFail; + } + if (scanRet != ErrorCode::Success) return scanRet; #else while (unprocessed > 0) { @@ -456,9 +504,9 @@ namespace SPTAG if (p_exWorkSpace->Deduper().CheckAndSet(vectorID)) continue; (this->*m_parseEncoding)(listInfo, (ValueType*)(p_postingListFullData + offsetVector)); - auto distance2leaf = m_headIndex->ComputeDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); + auto distance2leaf = PostingDistance(queryResults.GetQuantizedTarget(), p_postingListFullData + offsetVector); p_results.emplace_back(vectorID, distance2leaf, ByteArray::c_empty, - queryResults.WithVec() ? ByteArray::Alloc((std::uint8_t*)(p_postingListFullData + offsetVector), sizeof(ValueType) * m_opt->m_dim) : ByteArray::c_empty); + queryResults.WithVec() ? ByteArray::Alloc((std::uint8_t*)(p_postingListFullData + offsetVector), StoredVectorBytes()) : ByteArray::c_empty); } return ErrorCode::Success; }; @@ -474,7 +522,7 @@ namespace SPTAG auto curPostingID = p_exWorkSpace->m_postingIDs[pi]; auto it = m_globalVectorIDToHeadMap.find(curPostingID); if (it == m_globalVectorIDToHeadMap.end()) { - auto& request = p_exWorkSpace->m_diskRequests[pi]; + auto& request = DiskRequest(p_exWorkSpace, pi); request.m_readSize = 0; request.m_success = false; request.m_callback = nullptr; @@ -491,7 +539,7 @@ namespace SPTAG size_t totalBytes = (static_cast(listInfo->listPageCount) << PageSizeEx); #ifdef ASYNC_READ - auto& request = p_exWorkSpace->m_diskRequests[pi]; + auto& request = DiskRequest(p_exWorkSpace, pi); request.m_offset = listInfo->listOffset; request.m_readSize = totalBytes; request.m_status = (fileid << 16) | (request.m_status & 0xffff); @@ -499,20 +547,22 @@ namespace SPTAG request.m_success = false; #ifdef BATCH_READ - request.m_callback = [&appendPosting, &request, &scanRet](bool success) + Helper::AsyncReadRequest* requestPtr = &request; + request.m_callback = [&appendPosting, requestPtr, &scanRet](bool success) { if (!success) { scanRet = ErrorCode::DiskIOFail; return; } - ErrorCode ret = appendPosting(request.m_buffer, static_cast(request.m_payload)); + ErrorCode ret = appendPosting(requestPtr->m_buffer, static_cast(requestPtr->m_payload)); if (ret != ErrorCode::Success) scanRet = ret; }; #else - request.m_callback = [&p_exWorkSpace, &request](bool success) + Helper::AsyncReadRequest* requestPtr = &request; + request.m_callback = [p_exWorkSpace, requestPtr](bool success) { - p_exWorkSpace->m_processIocp.push(&request); + p_exWorkSpace->m_processIocp.push(requestPtr); }; ++unprocessed; @@ -592,7 +642,7 @@ namespace SPTAG size_t totalBytes = (static_cast(listInfo->listPageCount) << PageSizeEx); #ifdef ASYNC_READ - auto& request = p_exWorkSpace->m_diskRequests[pi]; + auto& request = DiskRequest(p_exWorkSpace, pi); request.m_offset = listInfo->listOffset; request.m_readSize = totalBytes; request.m_status = (fileid << 16) | (request.m_status & 0xffff); @@ -617,9 +667,10 @@ namespace SPTAG */ }; #else // async read - request.m_callback = [&p_exWorkSpace, &request](bool success) + Helper::AsyncReadRequest* requestPtr = &request; + request.m_callback = [p_exWorkSpace, requestPtr](bool success) { - p_exWorkSpace->m_processIocp.push(&request); + p_exWorkSpace->m_processIocp.push(requestPtr); }; ++unprocessed; @@ -658,6 +709,7 @@ namespace SPTAG success = BatchReadFileAsync(m_indexFiles, (p_exWorkSpace->m_diskRequests).data(), postingListCount); retry++; } + if (!success) return ErrorCode::DiskIOFail; #else while (unprocessed > 0) { @@ -700,7 +752,8 @@ namespace SPTAG continue; } char* buffer = (char*)((p_exWorkSpace->m_pageBuffers[p_exWorkSpace->m_pi]).GetBuffer()); - ListInfo* listInfo = static_cast(p_exWorkSpace->m_diskRequests[p_exWorkSpace->m_pi].m_payload); + ListInfo* listInfo = static_cast( + DiskRequest(p_exWorkSpace, p_exWorkSpace->m_pi).m_payload); // decompress posting list char* p_postingListFullData = buffer + listInfo->pageOffset; if (m_enableDataCompression && p_exWorkSpace->m_offset == 0) @@ -744,6 +797,7 @@ namespace SPTAG size_t p_postingListSize, Selection &p_selections, std::shared_ptr p_fullVectors, + std::shared_ptr p_quantizedVectors, COMMON::Dataset& p_localToGlobal, bool p_enableDeltaEncoding = false, bool p_enablePostingListRearrange = false, @@ -769,7 +823,29 @@ namespace SPTAG vectorID.append(reinterpret_cast(&vid), sizeof(SizeType)); ValueType *p_vector = reinterpret_cast(p_fullVectors->GetVector(vid)); - if (p_enableDeltaEncoding) + const bool quantizePosting = + m_headIndex->m_pQuantizer && + p_fullVectors->GetValueType() == m_headIndex->m_pQuantizer->GetReconstructType() && + p_fullVectors->Dimension() == m_headIndex->m_pQuantizer->ReconstructDim(); + if (quantizePosting && p_quantizedVectors) + { + const void* quantizedVector = p_quantizedVectors->GetVector(vid); + vector.append( + reinterpret_cast(quantizedVector), + p_quantizedVectors->PerVectorDataSize()); + } + else if (quantizePosting) + { + thread_local std::vector quantizedVector; + quantizedVector.resize(static_cast( + m_headIndex->m_pQuantizer->GetNumSubvectors())); + m_headIndex->m_pQuantizer->QuantizeVector( + p_vector, quantizedVector.data(), false); + vector.append( + reinterpret_cast(quantizedVector.data()), + quantizedVector.size()); + } + else if (p_enableDeltaEncoding) { DimensionType n = p_fullVectors->Dimension(); std::vector p_vector_delta(n); @@ -831,10 +907,60 @@ namespace SPTAG SizeType fullCount = 0; size_t vectorInfoSize = 0; + std::shared_ptr quantizedVectors; { auto fullVectors = p_reader->GetVectorSet(); fullCount = fullVectors->Count(); - vectorInfoSize = fullVectors->PerVectorDataSize() + sizeof(SizeType); + const bool quantizePosting = + m_headIndex->m_pQuantizer && + fullVectors->GetValueType() == m_headIndex->m_pQuantizer->GetReconstructType() && + fullVectors->Dimension() == m_headIndex->m_pQuantizer->ReconstructDim(); + if (quantizePosting && p_opt.m_enableDeltaEncoding) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Delta encoding is not supported when SSD postings are quantized.\n"); + return false; + } + vectorInfoSize = (quantizePosting + ? static_cast(m_headIndex->m_pQuantizer->GetNumSubvectors()) + : fullVectors->PerVectorDataSize()) + sizeof(SizeType); + if (quantizePosting && !p_opt.m_quantizedVectorPath.empty()) + { + auto quantizedOptions = std::make_shared( + VectorValueType::UInt8, + m_headIndex->m_pQuantizer->GetNumSubvectors(), + VectorFileType::DEFAULT, + "|", + p_opt.m_iSSDNumberOfThreads, + false); + auto quantizedReader = Helper::VectorSetReader::CreateInstance(quantizedOptions); + if (!quantizedReader || + quantizedReader->LoadFile(p_opt.m_quantizedVectorPath) != ErrorCode::Success) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Failed to load pre-quantized posting vectors: %s\n", + p_opt.m_quantizedVectorPath.c_str()); + return false; + } + quantizedVectors = quantizedReader->GetVectorSet(); + if (!quantizedVectors || + quantizedVectors->GetValueType() != VectorValueType::UInt8 || + quantizedVectors->Dimension() != + m_headIndex->m_pQuantizer->GetNumSubvectors() || + quantizedVectors->Count() < fullCount || + quantizedVectors->PerVectorDataSize() != + m_headIndex->m_pQuantizer->GetNumSubvectors()) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Pre-quantized posting vectors are incompatible with the quantizer.\n"); + return false; + } + SPTAGLIB_LOG(Helper::LogLevel::LL_Info, + "Loaded pre-quantized posting vectors from %s (%d,%d).\n", + p_opt.m_quantizedVectorPath.c_str(), + quantizedVectors->Count(), + quantizedVectors->Dimension()); + } } if (upperBound > 0) fullCount = upperBound; @@ -1084,7 +1210,7 @@ namespace SPTAG headVector = (ValueType*)p_headIndex->GetSample(j); } std::string postingListFullData = GetPostingListFullData( - j, curPostingListSizes[j], selections, fullVectors, p_localToGlobal, p_opt.m_enableDeltaEncoding, p_opt.m_enablePostingListRearrange, headVector); + j, curPostingListSizes[j], selections, fullVectors, quantizedVectors, p_localToGlobal, p_opt.m_enableDeltaEncoding, p_opt.m_enablePostingListRearrange, headVector); samplesBuffer += postingListFullData; samplesSizes.push_back(postingListFullData.size()); @@ -1124,7 +1250,7 @@ namespace SPTAG } std::string postingListFullData = GetPostingListFullData(postingListId, postingListSize[postingListId], - selections, fullVectors, p_localToGlobal, p_opt.m_enableDeltaEncoding, + selections, fullVectors, quantizedVectors, p_localToGlobal, p_opt.m_enableDeltaEncoding, p_opt.m_enablePostingListRearrange, headVector); size_t sizeToCompress = postingListSize[postingListId] * vectorInfoSize; if (sizeToCompress != postingListFullData.size()) @@ -1189,7 +1315,7 @@ namespace SPTAG postPageNum, postPageOffset, postingOrderInIndex, - fullVectors, p_headToLocal, p_localToGlobal, + fullVectors, quantizedVectors, p_headToLocal, p_localToGlobal, curPostingListOffSet); } @@ -1499,7 +1625,10 @@ namespace SPTAG const std::unique_ptr& p_postPageNum, const std::unique_ptr& p_postPageOffset, const std::vector& p_postingOrderInIndex, - std::shared_ptr p_fullVectors, COMMON::Dataset& p_headToLocal, COMMON::Dataset& p_localToGlobal, + std::shared_ptr p_fullVectors, + std::shared_ptr p_quantizedVectors, + COMMON::Dataset& p_headToLocal, + COMMON::Dataset& p_localToGlobal, size_t p_postingListOffset) { SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Start output...\n"); @@ -1564,7 +1693,15 @@ namespace SPTAG } // Vector dimension - int i32Val = static_cast(p_fullVectors->Dimension()); + size_t storedVectorBytes = p_spacePerVector - sizeof(SizeType); + if (storedVectorBytes % sizeof(ValueType) != 0) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Posting vector bytes %zu are not aligned to value type size %zu.\n", + storedVectorBytes, sizeof(ValueType)); + throw std::runtime_error("Posting vector size is not value-type aligned"); + } + int i32Val = static_cast(storedVectorBytes / sizeof(ValueType)); if (ptr->WriteBinary(sizeof(i32Val), reinterpret_cast(&i32Val)) != sizeof(i32Val)) { SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Failed to write SSDIndex File!"); throw std::runtime_error("Failed to write SSDIndex File"); @@ -1710,7 +1847,7 @@ namespace SPTAG headVector = (ValueType *)p_headIndex->GetSample(postingListId); } std::string postingListFullData = GetPostingListFullData( - postingListId, p_postingListSizes[id], p_postingSelections, p_fullVectors, p_localToGlobal, p_enableDeltaEncoding, p_enablePostingListRearrange, headVector); + postingListId, p_postingListSizes[id], p_postingSelections, p_fullVectors, p_quantizedVectors, p_localToGlobal, p_enableDeltaEncoding, p_enablePostingListRearrange, headVector); size_t postingListFullSize = p_postingListSizes[id] * p_spacePerVector; if (postingListFullSize != postingListFullData.size()) { diff --git a/AnnService/inc/Core/SPANN/Index.h b/AnnService/inc/Core/SPANN/Index.h index 041acbf1b..dd2d9c90b 100644 --- a/AnnService/inc/Core/SPANN/Index.h +++ b/AnnService/inc/Core/SPANN/Index.h @@ -85,6 +85,11 @@ namespace SPTAG std::shared_ptr> m_freeWorkSpaceIds; std::atomic m_workspaceCount = 0; + bool UseQuantizerForIndexBuild() const + { + return m_pQuantizer && m_pQuantizer->QuantizeForIndexBuild(); + } + public: Index() { diff --git a/AnnService/inc/Core/SPANN/Options.h b/AnnService/inc/Core/SPANN/Options.h index 3de2d8fa2..a977278b7 100644 --- a/AnnService/inc/Core/SPANN/Options.h +++ b/AnnService/inc/Core/SPANN/Options.h @@ -44,6 +44,7 @@ namespace SPTAG { bool m_deleteHeadVectors; int m_ssdIndexFileNum; std::string m_quantizerFilePath; + std::string m_quantizedVectorPath; SizeType m_datasetRowsInBlock; SizeType m_datasetCapacity; diff --git a/AnnService/inc/Core/SPANN/ParameterDefinitionList.h b/AnnService/inc/Core/SPANN/ParameterDefinitionList.h index bd6a1f943..09fc5b234 100644 --- a/AnnService/inc/Core/SPANN/ParameterDefinitionList.h +++ b/AnnService/inc/Core/SPANN/ParameterDefinitionList.h @@ -33,6 +33,7 @@ DefineBasicParameter(m_ssdIndex, std::string, std::string("SPTAGFullList.bin"), DefineBasicParameter(m_deleteHeadVectors, bool, false, "DeleteHeadVectors") DefineBasicParameter(m_ssdIndexFileNum, int, 1, "SSDIndexFileNum") DefineBasicParameter(m_quantizerFilePath, std::string, std::string(), "QuantizerFilePath") +DefineBasicParameter(m_quantizedVectorPath, std::string, std::string(), "QuantizedVectorPath") DefineBasicParameter(m_datasetRowsInBlock, SizeType, 1024 * 1024, "DataBlockSize") DefineBasicParameter(m_datasetCapacity, SizeType, SPTAG::MaxSize, "DataCapacity") #endif diff --git a/AnnService/inc/SSDServing/SSDIndex.h b/AnnService/inc/SSDServing/SSDIndex.h index 8ecea88df..6c76fc8f9 100644 --- a/AnnService/inc/SSDServing/SSDIndex.h +++ b/AnnService/inc/SSDServing/SSDIndex.h @@ -181,11 +181,18 @@ namespace SPTAG { int K = p_opts.m_resultNum; int truthK = (p_opts.m_truthResultNum <= 0) ? K : p_opts.m_truthResultNum; ErrorCode ret; + const bool useADC = p_index->m_pQuantizer && p_opts.m_enableADC; + const VectorValueType queryValueType = useADC + ? p_index->m_pQuantizer->GetReconstructType() + : p_opts.m_valueType; + const DimensionType queryDimension = useADC + ? p_index->m_pQuantizer->ReconstructDim() + : p_opts.m_dim; if (!warmupFile.empty()) { SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Start loading warmup query set...\n"); - std::shared_ptr queryOptions(new Helper::ReaderOptions(p_opts.m_valueType, p_opts.m_dim, p_opts.m_warmupType, p_opts.m_warmupDelimiter)); + std::shared_ptr queryOptions(new Helper::ReaderOptions(queryValueType, queryDimension, p_opts.m_warmupType, p_opts.m_warmupDelimiter)); auto queryReader = Helper::VectorSetReader::CreateInstance(queryOptions); if (ErrorCode::Success != (ret = queryReader->LoadFile(p_opts.m_warmupPath))) { @@ -199,7 +206,15 @@ namespace SPTAG { std::vector warmpUpStats(warmupNumQueries); for (int i = 0; i < warmupNumQueries; ++i) { - (*((COMMON::QueryResultSet*)&warmupResults[i])).SetTarget(reinterpret_cast(warmupQuerySet->GetVector(i)), p_index->m_pQuantizer); + if (p_index->m_pQuantizer && p_index->m_pQuantizer->QuantizeForIndexBuild()) + { + (*((COMMON::QueryResultSet*)&warmupResults[i])) + .SetTarget(reinterpret_cast(warmupQuerySet->GetVector(i)), p_index->m_pQuantizer); + } + else + { + warmupResults[i].SetTarget(warmupQuerySet->GetVector(i)); + } warmupResults[i].Reset(); } @@ -209,7 +224,7 @@ namespace SPTAG { } SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Start loading QuerySet...\n"); - std::shared_ptr queryOptions(new Helper::ReaderOptions(p_opts.m_valueType, p_opts.m_dim, p_opts.m_queryType, p_opts.m_queryDelimiter)); + std::shared_ptr queryOptions(new Helper::ReaderOptions(queryValueType, queryDimension, p_opts.m_queryType, p_opts.m_queryDelimiter)); auto queryReader = Helper::VectorSetReader::CreateInstance(queryOptions); if (ErrorCode::Success != (ret = queryReader->LoadFile(p_opts.m_queryPath))) { @@ -223,7 +238,15 @@ namespace SPTAG { std::vector stats(numQueries); for (int i = 0; i < numQueries; ++i) { - (*((COMMON::QueryResultSet*)&results[i])).SetTarget(reinterpret_cast(querySet->GetVector(i)), p_index->m_pQuantizer); + if (p_index->m_pQuantizer && p_index->m_pQuantizer->QuantizeForIndexBuild()) + { + (*((COMMON::QueryResultSet*)&results[i])) + .SetTarget(reinterpret_cast(querySet->GetVector(i)), p_index->m_pQuantizer); + } + else + { + results[i].SetTarget(querySet->GetVector(i)); + } results[i].Reset(); } diff --git a/AnnService/src/Core/Common/RaBitQAutoTuner.cpp b/AnnService/src/Core/Common/RaBitQAutoTuner.cpp new file mode 100644 index 000000000..2d47457a9 --- /dev/null +++ b/AnnService/src/Core/Common/RaBitQAutoTuner.cpp @@ -0,0 +1,689 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "inc/Core/Common/RaBitQAutoTuner.h" + +#include "inc/Core/VectorIndex.h" +#include "inc/Helper/StringConvert.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace SPTAG +{ +namespace COMMON +{ +namespace +{ + +constexpr SizeType kBatchSize = 64 * 1024; +constexpr const char* kSection = "RaBitQAutoTune"; + +template +bool ReadRequired(const Helper::IniReader& p_config, + const char* p_section, + const char* p_name, + T& p_value, + std::string& p_error) +{ + if (!p_config.DoesParameterExist(p_section, p_name)) { + p_error = std::string("[") + p_section + "] " + p_name + " is required"; + return false; + } + const std::string raw = p_config.GetParameter( + p_section, p_name, std::string()); + if (!Helper::Convert::ConvertStringTo(raw.c_str(), p_value)) { + p_error = std::string("invalid [") + p_section + "] " + p_name + ": " + raw; + return false; + } + return true; +} + +bool ReadRequiredString(const Helper::IniReader& p_config, + const char* p_section, + const char* p_name, + std::string& p_value, + std::string& p_error) +{ + if (!p_config.DoesParameterExist(p_section, p_name) || + (p_value = p_config.GetParameter(p_section, p_name, std::string())).empty()) { + p_error = std::string("[") + p_section + "] " + p_name + " is required"; + return false; + } + return true; +} + +ErrorCode LoadTruth(const std::string& p_path, + TruthFileType p_type, + SizeType p_queryCount, + std::vector>& p_truth, + std::string& p_error) +{ + auto input = f_createIO(); + if (!input || !input->Initialize(p_path.c_str(), std::ios::binary | std::ios::in)) { + p_error = "cannot open [Base] TruthPath: " + p_path; + return ErrorCode::FailedOpenFile; + } + + p_truth.clear(); + p_truth.reserve(static_cast(p_queryCount)); + if (p_type == TruthFileType::DEFAULT) { + std::int32_t rows = 0; + std::int32_t depth = 0; + if (input->ReadBinary(sizeof(rows), reinterpret_cast(&rows)) != sizeof(rows) || + input->ReadBinary(sizeof(depth), reinterpret_cast(&depth)) != sizeof(depth) || + rows < p_queryCount || depth <= 0) { + p_error = "invalid or insufficient DEFAULT truth header"; + return ErrorCode::FailedParseValue; + } + for (SizeType query = 0; query < p_queryCount; ++query) { + std::vector ids(static_cast(depth)); + const std::uint64_t bytes = + sizeof(std::int32_t) * static_cast(depth); + if (input->ReadBinary(bytes, reinterpret_cast(ids.data())) != bytes) { + p_error = "DEFAULT truth ended before QueryCountLimit rows"; + return ErrorCode::FailedParseValue; + } + p_truth.emplace_back(ids.begin(), ids.end()); + } + } else if (p_type == TruthFileType::XVEC) { + std::int32_t expectedDepth = -1; + for (SizeType query = 0; query < p_queryCount; ++query) { + std::int32_t depth = 0; + if (input->ReadBinary(sizeof(depth), reinterpret_cast(&depth)) != sizeof(depth) || + depth <= 0 || (expectedDepth >= 0 && depth != expectedDepth)) { + p_error = "XVEC truth has missing or inconsistent candidate depth"; + return ErrorCode::FailedParseValue; + } + expectedDepth = depth; + std::vector ids(static_cast(depth)); + const std::uint64_t bytes = + sizeof(std::int32_t) * static_cast(depth); + if (input->ReadBinary(bytes, reinterpret_cast(ids.data())) != bytes) { + p_error = "XVEC truth ended before QueryCountLimit rows"; + return ErrorCode::FailedParseValue; + } + p_truth.emplace_back(ids.begin(), ids.end()); + } + } else if (p_type == TruthFileType::TXT) { + std::size_t expectedDepth = 0; + std::uint64_t bufferSize = 64 * 1024; + std::unique_ptr buffer(new char[bufferSize]); + for (SizeType query = 0; query < p_queryCount; ++query) { + if (input->ReadString(bufferSize, buffer, '\n') == 0) { + p_error = "TXT truth ended before QueryCountLimit rows"; + return ErrorCode::FailedParseValue; + } + std::vector row; + char* context = nullptr; +#ifdef _MSC_VER + char* token = strtok_s(buffer.get(), " \t", &context); +#else + char* token = strtok_r(buffer.get(), " \t", &context); +#endif + while (token != nullptr) { + SizeType id = -1; + if (!Helper::Convert::ConvertStringTo(token, id)) { + p_error = "TXT truth contains a non-integer candidate ID"; + return ErrorCode::FailedParseValue; + } + row.push_back(id); +#ifdef _MSC_VER + token = strtok_s(nullptr, " \t", &context); +#else + token = strtok_r(nullptr, " \t", &context); +#endif + } + if (row.empty() || (!p_truth.empty() && row.size() != expectedDepth)) { + p_error = "TXT truth has empty or inconsistent candidate depth"; + return ErrorCode::FailedParseValue; + } + expectedDepth = row.size(); + p_truth.emplace_back(std::move(row)); + } + } else { + p_error = "[Base] TruthType must be DEFAULT, XVEC, or TXT"; + return ErrorCode::FailedParseValue; + } + return ErrorCode::Success; +} + +ErrorCode TrainCentroid(const std::shared_ptr& p_reader, + DimensionType p_dimension, + std::shared_ptr& p_model, + SizeType& p_count, + std::string& p_error) +{ + p_model = std::make_shared(p_dimension, 1, false); + if (p_model->BeginTraining() != ErrorCode::Success) { + p_error = "failed to initialize streaming RaBitQ centroid training"; + return ErrorCode::Fail; + } + + p_count = 0; + for (SizeType start = 0;; start += kBatchSize) { + const auto batch = p_reader->GetVectorSet(start, start + kBatchSize); + if (!batch || batch->Count() == 0) { + break; + } + if (batch->GetValueType() != VectorValueType::Float || + batch->Dimension() != p_dimension || + p_model->AddTrainingBatch(batch) != ErrorCode::Success) { + p_error = "base vector batch is incompatible with Float RaBitQ training"; + return ErrorCode::FailedParseValue; + } + p_count += batch->Count(); + if (batch->Count() < kBatchSize) { + break; + } + if (start > (std::numeric_limits::max)() - 2 * kBatchSize) { + p_error = "base vector count exceeds native SizeType"; + return ErrorCode::Fail; + } + } + if (p_count <= 0 || p_model->FinishTraining() != ErrorCode::Success) { + p_error = "base vector source is empty"; + return ErrorCode::EmptyData; + } + return ErrorCode::Success; +} + +ErrorCode EvaluateBits(const std::shared_ptr& p_quantizer, + const std::shared_ptr& p_baseReader, + const std::shared_ptr& p_queries, + const std::vector>& p_truth, + SizeType p_baseCount, + int p_resultCount, + int p_threads, + float& p_recall, + std::string& p_error) +{ + std::vector candidateIds; + for (const auto& row : p_truth) { + candidateIds.insert(candidateIds.end(), row.begin(), row.end()); + } + std::sort(candidateIds.begin(), candidateIds.end()); + candidateIds.erase(std::unique(candidateIds.begin(), candidateIds.end()), candidateIds.end()); + + const DimensionType codeDimension = p_quantizer->GetNumSubvectors(); + std::vector codes( + candidateIds.size() * static_cast(codeDimension)); + p_quantizer->SetEnableADC(true); + + std::size_t candidate = 0; + for (SizeType start = 0; start < p_baseCount; start += kBatchSize) { + const SizeType end = std::min(p_baseCount, start + kBatchSize); + const auto batch = p_baseReader->GetVectorSet(start, end); + if (!batch || batch->Count() != end - start) { + p_error = "base vector source changed while evaluating RaBitQ bits"; + return ErrorCode::Fail; + } + while (candidate < candidateIds.size() && candidateIds[candidate] < end) { + const SizeType id = candidateIds[candidate]; + p_quantizer->QuantizeVector( + batch->GetVector(id - start), + codes.data() + candidate * static_cast(codeDimension), + false); + ++candidate; + } + } + if (candidate != candidateIds.size()) { + p_error = "not all truth candidates could be encoded"; + return ErrorCode::Fail; + } + + std::unordered_map codeOffsets; + codeOffsets.reserve(candidateIds.size()); + for (std::size_t i = 0; i < candidateIds.size(); ++i) { + codeOffsets.emplace(candidateIds[i], i); + } + + const int queryBytes = p_quantizer->QuantizeSize(); + std::vector queryCodes( + static_cast(p_queries->Count()) * static_cast(queryBytes)); + for (SizeType query = 0; query < p_queries->Count(); ++query) { + p_quantizer->QuantizeVector( + p_queries->GetVector(query), + queryCodes.data() + static_cast(query) * queryBytes, + true); + } + + std::vector queryRecalls(static_cast(p_queries->Count()), 0.0F); + std::atomic nextQuery(0); + std::atomic evaluationFailed(false); + const int workerCount = std::max(1, std::min(p_threads, p_queries->Count())); + std::vector workers; + workers.reserve(static_cast(workerCount)); + for (int worker = 0; worker < workerCount; ++worker) { + workers.emplace_back([&]() { + for (;;) { + const SizeType query = nextQuery.fetch_add(1); + if (query >= p_queries->Count()) { + return; + } + std::vector> ranked; + ranked.reserve(p_truth[static_cast(query)].size()); + const std::uint8_t* queryCode = + queryCodes.data() + static_cast(query) * queryBytes; + for (SizeType id : p_truth[static_cast(query)]) { + const auto offset = codeOffsets.find(id); + if (offset == codeOffsets.end()) { + evaluationFailed.store(true); + return; + } + const std::uint8_t* code = + codes.data() + offset->second * static_cast(codeDimension); + ranked.emplace_back(p_quantizer->L2Distance(queryCode, code), id); + } + std::partial_sort( + ranked.begin(), ranked.begin() + p_resultCount, ranked.end(), + [](const auto& p_left, const auto& p_right) { + return p_left.first < p_right.first || + (p_left.first == p_right.first && p_left.second < p_right.second); + }); + std::vector ids; + ids.reserve(static_cast(p_resultCount)); + for (int rank = 0; rank < p_resultCount; ++rank) { + ids.push_back(ranked[static_cast(rank)].second); + } + queryRecalls[static_cast(query)] = + RaBitQAutoTuner::RecallAtK( + p_truth[static_cast(query)], ids, p_resultCount); + } + }); + } + for (auto& worker : workers) { + worker.join(); + } + if (evaluationFailed.load()) { + p_error = "truth candidate code lookup failed during parallel evaluation"; + return ErrorCode::Fail; + } + double recallSum = 0.0; + for (float recall : queryRecalls) { + recallSum += recall; + } + p_recall = static_cast(recallSum / p_queries->Count()); + return ErrorCode::Success; +} + +ErrorCode SaveArtifacts(const std::shared_ptr& p_quantizer, + const std::shared_ptr& p_reader, + SizeType p_count, + const std::string& p_outputFolder, + RaBitQAutoTuneResult& p_result, + std::string& p_error) +{ + namespace fs = std::filesystem; + std::error_code filesystemError; + fs::create_directories(p_outputFolder, filesystemError); + if (filesystemError) { + p_error = "cannot create RaBitQ artifact directory: " + filesystemError.message(); + return ErrorCode::FailedCreateFile; + } + + const fs::path folder(p_outputFolder); + const fs::path marker = folder / "rabitq_auto.incomplete"; + const fs::path modelTemporary = folder / "rabitq_auto_quantizer.bin.incomplete"; + const fs::path vectorTemporary = folder / "rabitq_auto_vectors.bin.incomplete"; + const fs::path modelFinal = folder / "rabitq_auto_quantizer.bin"; + const fs::path vectorFinal = folder / "rabitq_auto_vectors.bin"; + fs::remove(modelTemporary, filesystemError); + fs::remove(vectorTemporary, filesystemError); + { + auto markerOutput = f_createIO(); + if (!markerOutput || + !markerOutput->Initialize(marker.string().c_str(), std::ios::out | std::ios::binary) || + markerOutput->WriteString("RaBitQ auto-tuning artifacts are incomplete\n") == 0) { + p_error = "cannot create RaBitQ incomplete marker"; + return ErrorCode::FailedCreateFile; + } + } + + auto fail = [&](ErrorCode p_status, const std::string& p_message) { + p_error = p_message; + fs::remove(modelTemporary, filesystemError); + fs::remove(vectorTemporary, filesystemError); + return p_status; + }; + + { + auto modelOutput = f_createIO(); + if (!modelOutput || + !modelOutput->Initialize( + modelTemporary.string().c_str(), std::ios::out | std::ios::binary) || + p_quantizer->SaveQuantizer(modelOutput) != ErrorCode::Success) { + return fail(ErrorCode::DiskIOFail, "failed to write RaBitQ quantizer"); + } + } + + const DimensionType codeDimension = p_quantizer->GetNumSubvectors(); + { + auto vectorOutput = f_createIO(); + if (!vectorOutput || + !vectorOutput->Initialize( + vectorTemporary.string().c_str(), std::ios::out | std::ios::binary) || + vectorOutput->WriteBinary(sizeof(p_count), reinterpret_cast(&p_count)) != + sizeof(p_count) || + vectorOutput->WriteBinary( + sizeof(codeDimension), reinterpret_cast(&codeDimension)) != + sizeof(codeDimension)) { + return fail(ErrorCode::DiskIOFail, "failed to initialize encoded vector artifact"); + } + p_quantizer->SetEnableADC(true); + std::vector code(static_cast(codeDimension)); + SizeType written = 0; + for (SizeType start = 0; start < p_count; start += kBatchSize) { + const SizeType end = std::min(p_count, start + kBatchSize); + const auto batch = p_reader->GetVectorSet(start, end); + if (!batch || batch->Count() != end - start) { + return fail(ErrorCode::Fail, "base vector source changed during final encoding"); + } + for (SizeType i = 0; i < batch->Count(); ++i) { + p_quantizer->QuantizeVector(batch->GetVector(i), code.data(), false); + if (vectorOutput->WriteBinary( + code.size(), reinterpret_cast(code.data())) != code.size()) { + return fail(ErrorCode::DiskIOFail, "failed to stream encoded base vectors"); + } + ++written; + } + } + if (written != p_count) { + return fail(ErrorCode::Fail, "encoded vector count mismatch"); + } + } + + const std::uintmax_t expectedVectorSize = + sizeof(SizeType) + sizeof(DimensionType) + + static_cast(p_count) * static_cast(codeDimension); + if (fs::file_size(vectorTemporary, filesystemError) != expectedVectorSize || + filesystemError) { + return fail(ErrorCode::Fail, "encoded vector artifact size mismatch"); + } + + auto modelInput = f_createIO(); + if (!modelInput || + !modelInput->Initialize( + modelTemporary.string().c_str(), std::ios::in | std::ios::binary)) { + return fail(ErrorCode::FailedOpenFile, "cannot reopen generated RaBitQ quantizer"); + } + const auto loaded = IQuantizer::LoadIQuantizer(modelInput); + const auto loadedRaBitQ = std::dynamic_pointer_cast(loaded); + if (!loadedRaBitQ || loadedRaBitQ->Bits() != p_quantizer->Bits() || + loadedRaBitQ->Dimension() != p_quantizer->Dimension() || + loadedRaBitQ->GetNumSubvectors() != codeDimension) { + return fail(ErrorCode::FailedParseValue, "generated model is incompatible with encoded vectors"); + } + + fs::remove(modelFinal, filesystemError); + fs::rename(modelTemporary, modelFinal, filesystemError); + if (filesystemError) { + return fail(ErrorCode::DiskIOFail, "cannot publish RaBitQ quantizer: " + filesystemError.message()); + } + fs::remove(vectorFinal, filesystemError); + fs::rename(vectorTemporary, vectorFinal, filesystemError); + if (filesystemError) { + return fail(ErrorCode::DiskIOFail, "cannot publish encoded vectors: " + filesystemError.message()); + } + fs::remove(marker, filesystemError); + + p_result.quantizerPath = modelFinal.string(); + p_result.vectorPath = vectorFinal.string(); + p_result.codeDimension = codeDimension; + p_result.vectorCount = p_count; + p_result.quantizer = p_quantizer; + return ErrorCode::Success; +} + +} // namespace + +bool RaBitQAutoTuner::IsEnabled(const Helper::IniReader& p_config) +{ + return p_config.DoesSectionExist(kSection) && + p_config.GetParameter(kSection, "isExecute", false); +} + +ErrorCode RaBitQAutoTuner::Run(Helper::IniReader& p_config, + const std::string& p_outputFolder, + RaBitQAutoTuneResult& p_result, + std::string& p_error) +{ + p_result = RaBitQAutoTuneResult(); + p_error.clear(); + for (const auto& parameter : p_config.GetParameters(kSection)) { + if (!Helper::StrUtils::StrEqualIgnoreCase(parameter.first.c_str(), "isExecute") && + !Helper::StrUtils::StrEqualIgnoreCase(parameter.first.c_str(), "TargetRecall")) { + p_error = "unsupported [RaBitQAutoTune] parameter: " + parameter.first; + return ErrorCode::FailedParseValue; + } + } + + VectorValueType valueType = VectorValueType::Undefined; + DistCalcMethod distance = DistCalcMethod::Undefined; + DimensionType dimension = 0; + VectorFileType vectorType = VectorFileType::Undefined; + VectorFileType queryType = VectorFileType::Undefined; + TruthFileType truthType = TruthFileType::Undefined; + SizeType queryCount = 0; + int resultCount = 0; + int threads = 0; + float targetRecall = 0.0F; + std::string vectorPath; + std::string queryPath; + std::string truthPath; + if (!ReadRequired(p_config, "Base", "ValueType", valueType, p_error) || + !ReadRequired(p_config, "Base", "DistCalcMethod", distance, p_error) || + !ReadRequired(p_config, "Base", "Dim", dimension, p_error) || + !ReadRequired(p_config, "Base", "VectorType", vectorType, p_error) || + !ReadRequired(p_config, "Base", "QueryType", queryType, p_error) || + !ReadRequired(p_config, "Base", "TruthType", truthType, p_error) || + !ReadRequiredString(p_config, "Base", "VectorPath", vectorPath, p_error) || + !ReadRequiredString(p_config, "Base", "QueryPath", queryPath, p_error) || + !ReadRequiredString(p_config, "Base", "TruthPath", truthPath, p_error) || + !ReadRequired(p_config, "SearchSSDIndex", "QueryCountLimit", queryCount, p_error) || + !ReadRequired(p_config, "SearchSSDIndex", "ResultNum", resultCount, p_error) || + !ReadRequired(p_config, "BuildSSDIndex", "NumberOfThreads", threads, p_error) || + !ReadRequired(p_config, kSection, "TargetRecall", targetRecall, p_error)) { + return ErrorCode::FailedParseValue; + } + if (valueType != VectorValueType::Float || distance != DistCalcMethod::L2) { + p_error = "global RaBitQ auto-tuning requires [Base] ValueType=Float and DistCalcMethod=L2"; + return ErrorCode::FailedParseValue; + } + if (dimension <= 0 || queryCount <= 0 || resultCount <= 0 || threads <= 0 || + !std::isfinite(targetRecall) || targetRecall < 0.0F || targetRecall > 1.0F || + vectorType == VectorFileType::Undefined || queryType == VectorFileType::Undefined || + truthType == TruthFileType::Undefined || p_outputFolder.empty()) { + p_error = "invalid RaBitQ input dimension, counts, types, threads, target recall, or output folder"; + return ErrorCode::FailedParseValue; + } + + const std::string vectorDelimiter = + p_config.GetParameter("Base", "VectorDelimiter", std::string("|")); + const std::string queryDelimiter = + p_config.GetParameter("Base", "QueryDelimiter", std::string("|")); + auto baseOptions = std::make_shared( + VectorValueType::Float, dimension, vectorType, vectorDelimiter, threads, false); + auto baseReader = Helper::VectorSetReader::CreateInstance(baseOptions); + if (!baseReader || baseReader->LoadFile(vectorPath) != ErrorCode::Success) { + p_error = "failed to load [Base] VectorPath with its declared VectorType"; + return ErrorCode::FailedOpenFile; + } + + std::shared_ptr centroidModel; + SizeType baseCount = 0; + ErrorCode status = + TrainCentroid(baseReader, dimension, centroidModel, baseCount, p_error); + if (status != ErrorCode::Success) { + return status; + } + + auto queryOptions = std::make_shared( + VectorValueType::Float, dimension, queryType, queryDelimiter, threads, false); + auto queryReader = Helper::VectorSetReader::CreateInstance(queryOptions); + if (!queryReader || queryReader->LoadFile(queryPath) != ErrorCode::Success) { + p_error = "failed to load [Base] QueryPath with its declared QueryType"; + return ErrorCode::FailedOpenFile; + } + const auto queries = queryReader->GetVectorSet(0, queryCount); + if (!queries || queries->Count() != queryCount || + queries->GetValueType() != VectorValueType::Float || + queries->Dimension() != dimension) { + p_error = "QueryPath does not contain exactly QueryCountLimit usable Float queries"; + return ErrorCode::FailedParseValue; + } + + std::vector> truth; + status = LoadTruth(truthPath, truthType, queryCount, truth, p_error); + if (status != ErrorCode::Success || + (status = ValidateTruth( + truth, baseCount, queryCount, resultCount, p_error)) != ErrorCode::Success) { + return status; + } + + std::shared_ptr selected; + status = SelectMinimumBits( + targetRecall, + [&](int p_bits, float& p_recall) { + auto candidate = centroidModel->CreateWithBits(p_bits); + if (!candidate) { + p_error = "failed to create RaBitQ candidate from shared centroid"; + return ErrorCode::Fail; + } + const ErrorCode evaluation = EvaluateBits( + candidate, baseReader, queries, truth, baseCount, + resultCount, threads, p_recall, p_error); + SPTAGLIB_LOG( + Helper::LogLevel::LL_Info, + "RaBitQ auto-tuning bits=%d Recall@%d=%.6f target=%.6f\n", + p_bits, resultCount, p_recall, targetRecall); + if (evaluation == ErrorCode::Success && p_recall >= targetRecall) { + selected = std::move(candidate); + } + return evaluation; + }, + p_result.selectedBits, p_result.recall); + if (status != ErrorCode::Success) { + if (p_error.empty()) { + p_error = "no RaBitQ bit width in the fixed range 1..8 meets TargetRecall"; + } + return status; + } + if (!selected || selected->Bits() != p_result.selectedBits) { + p_error = "RaBitQ selected model does not match selected bit width"; + return ErrorCode::Fail; + } + + status = SaveArtifacts( + selected, baseReader, baseCount, p_outputFolder, p_result, p_error); + if (status == ErrorCode::Success) { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Info, + "RaBitQ auto-tuning selected %d bits (Recall@%d=%.6f); encoded %d vectors at width %d\n", + p_result.selectedBits, resultCount, p_result.recall, + p_result.vectorCount, p_result.codeDimension); + } + return status; +} + +ErrorCode RaBitQAutoTuner::SelectMinimumBits(float p_targetRecall, + const BitEvaluator& p_evaluator, + int& p_selectedBits, + float& p_selectedRecall) +{ + p_selectedBits = 0; + p_selectedRecall = 0.0F; + if (!p_evaluator || !std::isfinite(p_targetRecall) || + p_targetRecall < 0.0F || p_targetRecall > 1.0F) { + return ErrorCode::FailedParseValue; + } + int low = 1; + int high = 8; + while (low <= high) { + const int bits = low + (high - low) / 2; + float recall = 0.0F; + const ErrorCode status = p_evaluator(bits, recall); + if (status != ErrorCode::Success || !std::isfinite(recall)) { + return status == ErrorCode::Success ? ErrorCode::Fail : status; + } + if (recall >= p_targetRecall) { + p_selectedBits = bits; + p_selectedRecall = recall; + high = bits - 1; + } + else { + low = bits + 1; + } + } + return p_selectedBits == 0 ? ErrorCode::Fail : ErrorCode::Success; +} + +ErrorCode RaBitQAutoTuner::ValidateTruth( + const std::vector>& p_truth, + SizeType p_baseCount, + SizeType p_queryCount, + int p_resultCount, + std::string& p_error) +{ + if (p_baseCount <= 0 || p_queryCount <= 0 || p_resultCount <= 0 || + p_truth.size() != static_cast(p_queryCount)) { + p_error = "truth query count does not match QueryCountLimit"; + return ErrorCode::FailedParseValue; + } + std::size_t depth = 0; + for (std::size_t query = 0; query < p_truth.size(); ++query) { + const auto& row = p_truth[query]; + if (query == 0) { + depth = row.size(); + if (depth <= static_cast(p_resultCount)) { + p_error = "truth candidate depth must be greater than ResultNum"; + return ErrorCode::FailedParseValue; + } + } else if (row.size() != depth) { + p_error = "truth candidate depth is inconsistent across queries"; + return ErrorCode::FailedParseValue; + } + std::unordered_set seen; + seen.reserve(row.size()); + for (SizeType id : row) { + if (id < 0 || id >= p_baseCount) { + p_error = "truth contains a candidate ID outside the base vector source"; + return ErrorCode::FailedParseValue; + } + if (!seen.insert(id).second) { + p_error = "truth contains a duplicate candidate ID"; + return ErrorCode::FailedParseValue; + } + } + } + return ErrorCode::Success; +} + +float RaBitQAutoTuner::RecallAtK(const std::vector& p_exact, + const std::vector& p_reranked, + int p_resultCount) +{ + if (p_resultCount <= 0 || + p_exact.size() < static_cast(p_resultCount) || + p_reranked.size() < static_cast(p_resultCount)) { + return 0.0F; + } + std::unordered_set exact( + p_exact.begin(), p_exact.begin() + p_resultCount); + int matches = 0; + for (int i = 0; i < p_resultCount; ++i) { + matches += exact.find(p_reranked[static_cast(i)]) != exact.end(); + } + return static_cast(matches) / p_resultCount; +} + +} // namespace COMMON +} // namespace SPTAG diff --git a/AnnService/src/Core/Common/RaBitQQuantizer.cpp b/AnnService/src/Core/Common/RaBitQQuantizer.cpp index 82ef4ee8f..446172c72 100644 --- a/AnnService/src/Core/Common/RaBitQQuantizer.cpp +++ b/AnnService/src/Core/Common/RaBitQQuantizer.cpp @@ -41,34 +41,83 @@ ErrorCode RaBitQQuantizer::Initialize(DimensionType p_dimension, int p_bits, boo m_quantizer_config = rabitqlib::quant::faster_config( static_cast(m_padded_dimension), static_cast(m_bits)); m_ip_func = rabitqlib::select_excode_ipfunc(static_cast(m_bits)); + m_training_sum.clear(); + m_training_count = 0; + m_trained = false; return ErrorCode::Success; } ErrorCode RaBitQQuantizer::Train(const std::shared_ptr& p_vectors) { - if (!Ready() || !p_vectors || p_vectors->GetValueType() != VectorValueType::Float || - p_vectors->Dimension() != m_dimension || p_vectors->Count() <= 0) { + ErrorCode status = BeginTraining(); + if (status == ErrorCode::Success) { + status = AddTrainingBatch(p_vectors); + } + return status == ErrorCode::Success ? FinishTraining() : status; +} + +ErrorCode RaBitQQuantizer::BeginTraining() +{ + if (!Ready()) { return ErrorCode::FailedParseValue; } + m_training_sum.assign(static_cast(m_dimension), 0.0); + m_training_count = 0; + m_trained = false; + return ErrorCode::Success; +} - std::vector accumulator(static_cast(m_dimension), 0.0); +ErrorCode RaBitQQuantizer::AddTrainingBatch(const std::shared_ptr& p_vectors) +{ + if (!Ready() || m_training_sum.size() != static_cast(m_dimension) || + !p_vectors || p_vectors->GetValueType() != VectorValueType::Float || + p_vectors->Dimension() != m_dimension || p_vectors->Count() <= 0) { + return ErrorCode::FailedParseValue; + } + const auto batchCount = static_cast(p_vectors->Count()); + if (m_training_count > (std::numeric_limits::max)() - batchCount) { + return ErrorCode::Fail; + } std::vector prepared; for (SizeType i = 0; i < p_vectors->Count(); ++i) { const auto* vector = static_cast(p_vectors->GetVector(i)); PrepareInput(vector, prepared); for (DimensionType j = 0; j < m_dimension; ++j) { - accumulator[static_cast(j)] += prepared[static_cast(j)]; + m_training_sum[static_cast(j)] += prepared[static_cast(j)]; } } + m_training_count += batchCount; + return ErrorCode::Success; +} - const double inverse_count = 1.0 / static_cast(p_vectors->Count()); +ErrorCode RaBitQQuantizer::FinishTraining() +{ + if (!Ready() || m_training_count == 0 || + m_training_sum.size() != static_cast(m_dimension)) { + return ErrorCode::FailedParseValue; + } + const double inverse_count = 1.0 / static_cast(m_training_count); for (DimensionType j = 0; j < m_dimension; ++j) { m_centroid[static_cast(j)] = - static_cast(accumulator[static_cast(j)] * inverse_count); + static_cast(m_training_sum[static_cast(j)] * inverse_count); } + m_training_sum.clear(); + m_training_count = 0; + m_trained = true; return ErrorCode::Success; } +std::shared_ptr RaBitQQuantizer::CreateWithBits(int p_bits) const +{ + if (!Ready() || !m_trained || p_bits < 1 || p_bits > 8) { + return nullptr; + } + auto quantizer = std::make_shared(m_dimension, p_bits, m_normalize); + quantizer->m_centroid = m_centroid; + quantizer->m_trained = true; + return quantizer; +} + float RaBitQQuantizer::L2Distance(const std::uint8_t* p_x, const std::uint8_t* p_y) const { thread_local std::vector reconstructed_query; @@ -246,7 +295,7 @@ std::uint64_t RaBitQQuantizer::BufferSize() const ErrorCode RaBitQQuantizer::SaveQuantizer(std::shared_ptr p_output) const { - if (!p_output || !Ready()) { + if (!p_output || !Ready() || !m_trained) { return ErrorCode::Fail; } @@ -283,6 +332,7 @@ ErrorCode RaBitQQuantizer::LoadQuantizer(std::shared_ptr p_input m_centroid.size() * sizeof(float)) { return ErrorCode::FailedParseValue; } + m_trained = true; return ErrorCode::Success; } @@ -299,6 +349,7 @@ ErrorCode RaBitQQuantizer::LoadQuantizer(std::uint8_t* p_raw_bytes) } p_raw_bytes += sizeof(header); std::memcpy(m_centroid.data(), p_raw_bytes, m_centroid.size() * sizeof(float)); + m_trained = true; return ErrorCode::Success; } diff --git a/AnnService/src/Core/SPANN/SPANNIndex.cpp b/AnnService/src/Core/SPANN/SPANNIndex.cpp index 35f7a72f5..bc94f3c36 100644 --- a/AnnService/src/Core/SPANN/SPANNIndex.cpp +++ b/AnnService/src/Core/SPANN/SPANNIndex.cpp @@ -57,10 +57,21 @@ template void Index::SetQuantizer(std::shared_ptrDistanceCalcSelector(m_options.m_distCalcMethod); - m_iBaseSquare = (m_options.m_distCalcMethod == DistCalcMethod::Cosine) - ? m_pQuantizer->GetBase() * m_pQuantizer->GetBase() - : 1; + m_pQuantizer->SetEnableADC(m_options.m_enableADC); + if (UseQuantizerForIndexBuild()) + { + m_fComputeDistance = m_pQuantizer->DistanceCalcSelector(m_options.m_distCalcMethod); + m_iBaseSquare = (m_options.m_distCalcMethod == DistCalcMethod::Cosine) + ? m_pQuantizer->GetBase() * m_pQuantizer->GetBase() + : 1; + } + else + { + m_fComputeDistance = COMMON::DistanceCalcSelector(m_options.m_distCalcMethod); + m_iBaseSquare = (m_options.m_distCalcMethod == DistCalcMethod::Cosine) + ? COMMON::Utils::GetBase() * COMMON::Utils::GetBase() + : 1; + } } else { @@ -103,7 +114,10 @@ template ErrorCode Index::LoadConfig(Helper::IniReader &p_reader template ErrorCode Index::LoadIndexDataFromMemory(const std::vector &p_indexBlobs) { /** Need to modify **/ - m_topIndex->SetQuantizer(m_pQuantizer); + if (UseQuantizerForIndexBuild()) + { + m_topIndex->SetQuantizer(m_pQuantizer); + } if (!m_options.m_persistentBufferPath.empty() && !direxists(m_options.m_persistentBufferPath.c_str())) mkdir(m_options.m_persistentBufferPath.c_str()); @@ -156,7 +170,10 @@ template ErrorCode Index::LoadIndexDataFromMemory(const std::vec template ErrorCode Index::LoadIndexData(const std::vector> &p_indexStreams) { - m_topIndex->SetQuantizer(m_pQuantizer); + if (UseQuantizerForIndexBuild()) + { + m_topIndex->SetQuantizer(m_pQuantizer); + } if (!m_options.m_persistentBufferPath.empty() && !direxists(m_options.m_persistentBufferPath.c_str())) mkdir(m_options.m_persistentBufferPath.c_str()); @@ -339,8 +356,14 @@ template ErrorCode Index::SearchIndex(QueryResult &p_query, Sear if (p_query.GetResultNum() >= m_options.m_searchInternalResultNum) p_queryResults = (COMMON::QueryResultSet *)&p_query; else + { p_queryResults = new COMMON::QueryResultSet((const T *)p_query.GetTarget(), m_options.m_searchInternalResultNum, p_query.WithMeta(), p_query.WithVec()); + if (UseQuantizerForIndexBuild()) + { + p_queryResults->SetTarget((const T *)p_query.GetTarget(), m_pQuantizer); + } + } ErrorCode ret; auto searchStart = std::chrono::high_resolution_clock::now(); @@ -667,8 +690,14 @@ ErrorCode Index::SearchHeadIndex(QueryResult& p_query, int p_tolayer, ExtraWo if (p_query.GetResultNum() >= m_options.m_searchInternalResultNum) p_queryResults = (COMMON::QueryResultSet *)&p_query; else + { p_queryResults = new COMMON::QueryResultSet((const T *)p_query.GetTarget(), m_options.m_searchInternalResultNum, p_query.WithMeta(), p_query.WithVec()); + if (UseQuantizerForIndexBuild()) + { + p_queryResults->SetTarget((const T *)p_query.GetTarget(), m_pQuantizer); + } + } ErrorCode ret; if ((ret = m_topIndex->SearchIndex(*p_queryResults)) != ErrorCode::Success) @@ -722,7 +751,7 @@ ErrorCode Index::SearchDiskIndex(QueryResult &p_query, SearchStats *p_stats, COMMON::QueryResultSet localResults((const T *)p_query.GetTarget(), m_options.m_searchInternalResultNum, p_query.WithMeta(), p_query.WithVec()); std::vector headCandidates; headCandidates.reserve(m_options.m_searchInternalResultNum); - if (m_pQuantizer && p_query.HasQuantizedTarget()) + if (m_pQuantizer) { localResults.SetTarget((const T *)p_query.GetTarget(), m_pQuantizer); } @@ -1118,7 +1147,9 @@ bool Index::SelectHeadInternal(std::shared_ptr &p_re bkt->m_iTreeNumber = m_options.m_iTreeNumber; bkt->m_fBalanceFactor = m_options.m_fBalanceFactor; bkt->m_parallelBuild = m_options.m_parallelBKTBuild; - bkt->m_pQuantizer = m_pQuantizer; + bkt->m_pQuantizer = UseQuantizerForIndexBuild() + ? m_pQuantizer + : nullptr; SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Start invoking BuildTrees.\n"); SPTAGLIB_LOG( Helper::LogLevel::LL_Info, @@ -1229,6 +1260,33 @@ bool Index::SelectHeadInternal(std::shared_ptr &p_re template ErrorCode Index::BuildIndexInternalLayer(std::shared_ptr &p_reader) { + struct ScopedIndexBuildADCMode + { + std::shared_ptr m_quantizer; + bool m_restore; + bool m_enableADC; + + ScopedIndexBuildADCMode(std::shared_ptr quantizer, bool quantizedIndexBuild) + : m_quantizer(std::move(quantizer)), + m_restore(m_quantizer != nullptr && quantizedIndexBuild), + m_enableADC(false) + { + if (m_restore) + { + m_enableADC = m_quantizer->GetEnableADC(); + m_quantizer->SetEnableADC(false); + } + } + + ~ScopedIndexBuildADCMode() + { + if (m_restore) + { + m_quantizer->SetEnableADC(m_enableADC); + } + } + } scopedADCMode(m_pQuantizer, UseQuantizerForIndexBuild()); + int currentLayer = static_cast(m_extraSearchers.size()); COMMON::Dataset localToGlobalID; { @@ -1258,7 +1316,7 @@ template ErrorCode Index::BuildIndexInternalLayer(std::shared_pt if (m_options.m_selectHead && m_topIndex == nullptr) { bool success = false; - if (m_pQuantizer) + if (UseQuantizerForIndexBuild()) { success = SelectHeadInternal(p_reader); } @@ -1279,13 +1337,21 @@ template ErrorCode Index::BuildIndexInternalLayer(std::shared_pt SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Begin Build Head...\n"); if (m_options.m_buildHead && m_topIndex == nullptr) { - auto valueType = m_pQuantizer ? SPTAG::VectorValueType::UInt8 : m_options.m_valueType; - auto dims = m_pQuantizer ? m_pQuantizer->GetNumSubvectors() : m_options.m_dim; + const bool quantizedIndexBuild = UseQuantizerForIndexBuild(); + auto valueType = quantizedIndexBuild + ? SPTAG::VectorValueType::UInt8 + : m_options.m_valueType; + auto dims = quantizedIndexBuild + ? m_pQuantizer->GetNumSubvectors() + : m_options.m_dim; m_topIndex = SPTAG::VectorIndex::CreateInstance(m_options.m_indexAlgoType, valueType); m_topIndex->SetParameter("DistCalcMethod", SPTAG::Helper::Convert::ConvertToString(m_options.m_distCalcMethod)); m_topIndex->SetParameter("ParallelBKTBuild", m_options.m_parallelBKTBuild ? "true" : "false"); - m_topIndex->SetQuantizer(m_pQuantizer); + if (quantizedIndexBuild) + { + m_topIndex->SetQuantizer(m_pQuantizer); + } for (const auto &iter : m_topParameters) { m_topIndex->SetParameter(iter.first.c_str(), iter.second.c_str()); @@ -1307,9 +1373,11 @@ template ErrorCode Index::BuildIndexInternalLayer(std::shared_pt SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Failed to build head index.\n"); return ErrorCode::Fail; } - if (!m_options.m_quantizerFilePath.empty()) + if (!m_options.m_quantizerFilePath.empty() && quantizedIndexBuild) + { m_topIndex->SetQuantizerFileName( m_options.m_quantizerFilePath.substr(m_options.m_quantizerFilePath.find_last_of("/\\") + 1)); + } if (m_topIndex->SaveIndex(m_options.m_indexDirectory + FolderSep + m_options.m_headIndexFolder) != ErrorCode::Success) { @@ -1340,7 +1408,10 @@ template ErrorCode Index::BuildIndexInternalLayer(std::shared_pt (m_options.m_indexDirectory + FolderSep + m_options.m_headIndexFolder).c_str()); return ErrorCode::Fail; } - m_topIndex->SetQuantizer(m_pQuantizer); + if (UseQuantizerForIndexBuild()) + { + m_topIndex->SetQuantizer(m_pQuantizer); + } if (!CheckHeadIndexType()) return ErrorCode::Fail; @@ -1474,7 +1545,10 @@ template ErrorCode Index::BuildIndexInternal(std::shared_ptrSetQuantizer(m_pQuantizer); + if (UseQuantizerForIndexBuild()) + { + m_topIndex->SetQuantizer(m_pQuantizer); + } m_topIndex->SetParameter("NumberOfThreads", std::to_string(m_options.m_iSSDNumberOfThreads)); m_topIndex->SetParameter("MaxCheck", std::to_string(m_options.m_maxCheck)); m_topIndex->SetParameter("HashTableExponent", std::to_string(m_options.m_hashExp)); @@ -1661,8 +1735,13 @@ template ErrorCode Index::BuildIndexInternal(std::shared_ptr ErrorCode Index::BuildIndex(bool p_normalized) { - SPTAG::VectorValueType valueType = m_pQuantizer ? SPTAG::VectorValueType::UInt8 : m_options.m_valueType; - SizeType dim = m_pQuantizer ? m_pQuantizer->GetNumSubvectors() : m_options.m_dim; + const bool quantizedIndexBuild = UseQuantizerForIndexBuild(); + SPTAG::VectorValueType valueType = quantizedIndexBuild + ? SPTAG::VectorValueType::UInt8 + : GetEnumValueType(); + SizeType dim = quantizedIndexBuild + ? m_pQuantizer->GetNumSubvectors() + : m_options.m_dim; std::shared_ptr vectorOptions( new Helper::ReaderOptions(valueType, dim, m_options.m_vectorType, m_options.m_vectorDelimiter, m_options.m_iSSDNumberOfThreads, p_normalized)); @@ -1709,7 +1788,10 @@ ErrorCode Index::BuildIndex(const void *p_data, SizeType p_vectorNum, Dimensi { vectorSet->Normalize(m_options.m_iSSDNumberOfThreads); } - SPTAG::VectorValueType valueType = m_pQuantizer ? SPTAG::VectorValueType::UInt8 : m_options.m_valueType; + SPTAG::VectorValueType valueType = + UseQuantizerForIndexBuild() + ? SPTAG::VectorValueType::UInt8 + : GetEnumValueType(); std::shared_ptr vectorOptions( new Helper::ReaderOptions(valueType, p_dimension, VectorFileType::DEFAULT, m_options.m_vectorDelimiter, m_options.m_iSSDNumberOfThreads, true)); @@ -1811,7 +1893,7 @@ template ErrorCode Index::SetParameter(const char *p_param, cons } if (SPTAG::Helper::StrUtils::StrEqualIgnoreCase(p_param, "DistCalcMethod")) { - if (m_pQuantizer) + if (UseQuantizerForIndexBuild()) { m_fComputeDistance = m_pQuantizer->DistanceCalcSelector(m_options.m_distCalcMethod); m_iBaseSquare = (m_options.m_distCalcMethod == DistCalcMethod::Cosine) @@ -1826,6 +1908,10 @@ template ErrorCode Index::SetParameter(const char *p_param, cons : 1; } } + if (SPTAG::Helper::StrUtils::StrEqualIgnoreCase(p_param, "EnableADC") && m_pQuantizer) + { + m_pQuantizer->SetEnableADC(m_options.m_enableADC); + } return ErrorCode::Success; } diff --git a/AnnService/src/Helper/AsyncFileReader.cpp b/AnnService/src/Helper/AsyncFileReader.cpp index 1245ab01f..2c1cf5249 100644 --- a/AnnService/src/Helper/AsyncFileReader.cpp +++ b/AnnService/src/Helper/AsyncFileReader.cpp @@ -126,7 +126,15 @@ bool BatchReadFileAsync(std::vector> &handlers, AsyncReadRequest *req = reinterpret_cast((events[i].data)); if (nullptr != req) { - req->m_callback(true); + const bool success = events[i].res == static_cast(req->m_readSize); + if (!success) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Async batch read failed at offset %llu, expected %zu bytes, actual %lld.\n", + req->m_offset, req->m_readSize, static_cast(events[i].res)); + } + req->m_success = success; + req->m_callback(success); } } totalQueued = totalDone; @@ -153,7 +161,15 @@ bool BatchReadFileAsync(std::vector> &handlers, AsyncReadRequest *req = reinterpret_cast((events[i].data)); if (nullptr != req) { - req->m_callback(true); + const bool success = events[i].res == static_cast(req->m_readSize); + if (!success) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "Async batch read failed at offset %llu, expected %zu bytes, actual %lld.\n", + req->m_offset, req->m_readSize, static_cast(events[i].res)); + } + req->m_success = success; + req->m_callback(success); } } return true; diff --git a/AnnService/src/Helper/VectorSetReader.cpp b/AnnService/src/Helper/VectorSetReader.cpp index 00089fe54..6cb752365 100644 --- a/AnnService/src/Helper/VectorSetReader.cpp +++ b/AnnService/src/Helper/VectorSetReader.cpp @@ -9,8 +9,9 @@ using namespace SPTAG; using namespace SPTAG::Helper; -ReaderOptions::ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, VectorFileType p_fileType, - std::string p_vectorDelimiter, std::uint32_t p_threadNum, bool p_normalized) +ReaderOptions::ReaderOptions(VectorValueType p_valueType, DimensionType p_dimension, + VectorFileType p_fileType, std::string p_vectorDelimiter, + std::uint32_t p_threadNum, bool p_normalized) : m_inputValueType(p_valueType), m_dimension(p_dimension), m_inputFileType(p_fileType), m_vectorDelimiter(p_vectorDelimiter), m_threadNum(p_threadNum), m_normalized(p_normalized) { diff --git a/AnnService/src/IndexBuilder/main.cpp b/AnnService/src/IndexBuilder/main.cpp index 56499840f..8f321e16a 100644 --- a/AnnService/src/IndexBuilder/main.cpp +++ b/AnnService/src/IndexBuilder/main.cpp @@ -7,14 +7,20 @@ #include "inc/Helper/VectorSetReader.h" #include +#include #include +#ifdef RABITQ +#include "inc/Core/Common/RaBitQAutoTuner.h" +#endif + using namespace SPTAG; class BuilderOptions : public Helper::ReaderOptions { public: - BuilderOptions() : Helper::ReaderOptions(VectorValueType::Float, 0, VectorFileType::TXT, "|", 32) + BuilderOptions() + : Helper::ReaderOptions(VectorValueType::Float, 0, VectorFileType::TXT, "|", 32) { AddRequiredOption(m_outputFolder, "-o", "--outputfolder", "Output folder."); AddRequiredOption(m_indexAlgoType, "-a", "--algo", "Index Algorithm type."); @@ -48,17 +54,6 @@ int main(int argc, char *argv[]) { exit(1); } - SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Set QuantizerFile = %s\n", options->m_quantizerFile.c_str()); - - auto indexBuilder = VectorIndex::CreateInstance(options->m_indexAlgoType, options->m_inputValueType); - if (!options->m_quantizerFile.empty()) - { - indexBuilder->LoadQuantizer(options->m_quantizerFile); - if (!indexBuilder->m_pQuantizer) - { - exit(1); - } - } Helper::IniReader iniReader; if (!options->m_builderConfigFile.empty() && @@ -89,6 +84,135 @@ int main(int argc, char *argv[]) paramVal.c_str()); } + std::string quantizerFile = options->m_quantizerFile; + VectorValueType builderValueType = options->m_inputValueType; + if (options->m_inputFiles.empty() && + iniReader.DoesParameterExist("Base", "ValueType")) + { + builderValueType = + iniReader.GetParameter("Base", "ValueType", VectorValueType::Undefined); + if (builderValueType == VectorValueType::Undefined) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, "Invalid [Base] ValueType.\n"); + return 1; + } + } + + const bool autoTuneEnabled = + iniReader.DoesSectionExist("RaBitQAutoTune") && + iniReader.GetParameter("RaBitQAutoTune", "isExecute", false); + if (autoTuneEnabled) + { +#ifdef RABITQ + if (options->m_indexAlgoType != IndexAlgoType::SPANN) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "RaBitQ auto-tuning is supported only for SPANN index construction.\n"); + return 1; + } + if (!options->m_inputFiles.empty()) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "RaBitQ auto-tuning reads vectors from the INI; --input is not allowed.\n"); + return 1; + } + + COMMON::RaBitQAutoTuneResult tuneResult; + std::string tuneError; + ErrorCode tuneStatus = ErrorCode::Fail; + try + { + tuneStatus = + COMMON::RaBitQAutoTuner::Run( + iniReader, options->m_outputFolder, tuneResult, tuneError); + } + catch (const std::exception& exception) + { + tuneError = exception.what(); + } + if (tuneStatus != ErrorCode::Success) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "RaBitQ auto-tuning failed: %s\n", tuneError.c_str()); + return 1; + } + + quantizerFile = tuneResult.quantizerPath; + iniReader.SetParameter("Base", "VectorSize", std::to_string(tuneResult.vectorCount)); + iniReader.SetParameter("Base", "QuantizerFilePath", tuneResult.quantizerPath); + iniReader.SetParameter("Base", "QuantizedVectorPath", tuneResult.vectorPath); + iniReader.SetParameter("BuildSSDIndex", "EnableADC", "true"); +#else + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "[RaBitQAutoTune] isExecute=true requires a build configured with RABITQ=ON.\n"); + return 1; +#endif + } + + if (quantizerFile.empty() && + iniReader.DoesParameterExist("Base", "QuantizerFilePath")) + { + quantizerFile = + iniReader.GetParameter("Base", "QuantizerFilePath", std::string()); + } + + SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Set QuantizerFile = %s\n", quantizerFile.c_str()); + + auto indexBuilder = VectorIndex::CreateInstance(options->m_indexAlgoType, builderValueType); + if (!indexBuilder) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Cannot create index builder.\n"); + return 1; + } + if (!quantizerFile.empty()) + { + if (indexBuilder->LoadQuantizer(quantizerFile) != ErrorCode::Success || + !indexBuilder->m_pQuantizer) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Cannot load quantizer file.\n"); + return 1; + } + if (!indexBuilder->m_pQuantizer->QuantizeForIndexBuild()) + { + const auto reconstructType = indexBuilder->m_pQuantizer->GetReconstructType(); + if (builderValueType != reconstructType) + { + if (!options->m_inputFiles.empty()) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "This quantizer requires raw reconstruct vectors for SPANN index build. " + "Set the input vector type to %s and keep pre-quantized codes in QuantizedVectorPath.\n", + Helper::Convert::ConvertToString(reconstructType).c_str()); + return 1; + } + if (iniReader.DoesParameterExist("Base", "VectorPath") && + !iniReader.DoesParameterExist("Base", "QuantizedVectorPath")) + { + SPTAGLIB_LOG( + Helper::LogLevel::LL_Error, + "This quantizer requires [Base] VectorPath to point to raw reconstruct vectors " + "and [Base] QuantizedVectorPath to point to pre-quantized codes.\n"); + return 1; + } + + builderValueType = reconstructType; + iniReader.SetParameter( + "Base", "ValueType", + Helper::Convert::ConvertToString(builderValueType)); + indexBuilder = VectorIndex::CreateInstance(options->m_indexAlgoType, builderValueType); + if (!indexBuilder || + indexBuilder->LoadQuantizer(quantizerFile) != ErrorCode::Success || + !indexBuilder->m_pQuantizer) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Cannot recreate index builder for quantizer reconstruct type.\n"); + return 1; + } + } + } + } + std::string sections[] = {"Base", "SelectHead", "BuildHead", "BuildSSDIndex", "Index"}; for (int i = 0; i < 5; i++) { @@ -106,6 +230,12 @@ int main(int argc, char *argv[]) std::shared_ptr vecset; if (options->m_inputFiles != "") { + if (options->m_dimension <= 0) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, + "--dimension is required when indexbuilder reads --input directly.\n"); + return 1; + } auto vectorReader = Helper::VectorSetReader::CreateInstance(options); if (ErrorCode::Success != vectorReader->LoadFile(options->m_inputFiles)) { @@ -118,16 +248,21 @@ int main(int argc, char *argv[]) } else { - if (!options->m_quantizerFile.empty()) + if (!quantizerFile.empty()) { indexBuilder->SetQuantizerFileName( - options->m_quantizerFile.substr(options->m_quantizerFile.find_last_of("/\\") + 1)); + quantizerFile.substr(quantizerFile.find_last_of("/\\") + 1)); } code = indexBuilder->BuildIndex(options->m_normalized); } if (code == ErrorCode::Success) { - indexBuilder->SaveIndex(options->m_outputFolder); + code = indexBuilder->SaveIndex(options->m_outputFolder); + if (code != ErrorCode::Success) + { + SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Failed to save index.\n"); + return 1; + } } else { diff --git a/AnnService/src/SSDServing/main.cpp b/AnnService/src/SSDServing/main.cpp index 867e02fea..64cd51bf2 100644 --- a/AnnService/src/SSDServing/main.cpp +++ b/AnnService/src/SSDServing/main.cpp @@ -136,14 +136,15 @@ int BootProgram(bool forANNIndexTestTool, std::mapm_generateTruth) { SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Start generating truth. It's maybe a long time.\n"); + VectorValueType vectorValueType = opts->m_valueType; SizeType dim = opts->m_dim; if (index->m_pQuantizer) { - valueType = VectorValueType::UInt8; - dim = index->m_pQuantizer->GetNumSubvectors(); + vectorValueType = index->m_pQuantizer->GetReconstructType(); + dim = index->m_pQuantizer->ReconstructDim(); } std::shared_ptr vectorOptions( - new Helper::ReaderOptions(valueType, dim, opts->m_vectorType, opts->m_vectorDelimiter)); + new Helper::ReaderOptions(vectorValueType, dim, opts->m_vectorType, opts->m_vectorDelimiter)); auto vectorReader = Helper::VectorSetReader::CreateInstance(vectorOptions); if (ErrorCode::Success != vectorReader->LoadFile(opts->m_vectorPath)) { diff --git a/Test/src/RaBitQQuantizerTest.cpp b/Test/src/RaBitQQuantizerTest.cpp index 56bd53b96..a1dc78c96 100644 --- a/Test/src/RaBitQQuantizerTest.cpp +++ b/Test/src/RaBitQQuantizerTest.cpp @@ -4,6 +4,7 @@ #include "inc/Test.h" #include "inc/Core/Common/QueryResultSet.h" +#include "inc/Core/Common/RaBitQAutoTuner.h" #include "inc/Core/Common/RaBitQQuantizer.h" #include "inc/Core/SPANN/Index.h" #include "inc/Core/VectorIndex.h" @@ -193,14 +194,14 @@ void VerifySpannSearch( std::filesystem::remove_all(index_directory); p_quantizer->SetEnableADC(false); - auto index = VectorIndex::CreateInstance(IndexAlgoType::SPANN, VectorValueType::UInt8); + auto index = VectorIndex::CreateInstance(IndexAlgoType::SPANN, VectorValueType::Float); BOOST_REQUIRE(index != nullptr); index->SetQuantizer(p_quantizer); ConfigureSpannIndex(index, index_directory, nullptr, p_storage, p_enable_compression); - BOOST_REQUIRE(index->BuildIndex(p_codes, nullptr, false, true) == ErrorCode::Success); + BOOST_REQUIRE(index->BuildIndex(p_raw, nullptr, false, true) == ErrorCode::Success); p_quantizer->SetEnableADC(true); - auto* spann_index = static_cast*>(index.get()); + auto* spann_index = static_cast*>(index.get()); std::vector head_ids; BOOST_REQUIRE(spann_index->GetHeadIndexMapping(1, head_ids) == ErrorCode::Success); SizeType expected = 0; @@ -212,16 +213,25 @@ void VerifySpannSearch( COMMON::QueryResultSet query( reinterpret_cast(p_raw->GetVector(expected)), 96); BOOST_REQUIRE(index->SearchIndex(query) == ErrorCode::Success); - - bool found = false; for (int rank = 0; rank < query.GetResultNum(); ++rank) { const auto* result = query.GetResult(rank); - if (result != nullptr && result->VID == expected) { - found = true; - break; + if (result != nullptr && result->VID != -1) { + BOOST_CHECK(std::isfinite(result->Dist)); } } - BOOST_CHECK(found); + + COMMON::QueryResultSet direct_query( + reinterpret_cast(p_raw->GetVector(expected)), 1); + direct_query.SetTarget( + reinterpret_cast(p_raw->GetVector(expected)), p_quantizer); + const auto* query_code = reinterpret_cast( + direct_query.GetQuantizedTarget()); + const auto* own_code = reinterpret_cast( + p_codes->GetVector(expected)); + const auto* far_code = reinterpret_cast( + p_codes->GetVector(kVectorCount - 1)); + BOOST_CHECK(p_quantizer->L2Distance(query_code, own_code) < + p_quantizer->L2Distance(query_code, far_code)); index.reset(); std::filesystem::remove_all(index_directory); @@ -247,7 +257,7 @@ void VerifySSDServingSearch( } p_quantizer->SetEnableADC(false); - auto index = VectorIndex::CreateInstance(IndexAlgoType::SPANN, VectorValueType::UInt8); + auto index = VectorIndex::CreateInstance(IndexAlgoType::SPANN, VectorValueType::Float); BOOST_REQUIRE(index != nullptr); index->SetQuantizer(p_quantizer); ConfigureSpannIndex(index, index_directory, kQueryFile, "FILEIO", false); @@ -256,9 +266,9 @@ void VerifySSDServingSearch( index->SetParameter("SearchInternalResultNum", "96", "SearchSSDIndex"); index->SetParameter("ResultNum", "8", "SearchSSDIndex"); index->SetParameter("QueryCountLimit", std::to_string(kSearchQueryCount), "SearchSSDIndex"); - BOOST_REQUIRE(index->BuildIndex(p_codes, nullptr, false, true) == ErrorCode::Success); + BOOST_REQUIRE(index->BuildIndex(p_raw, nullptr, false, true) == ErrorCode::Success); - auto* spann_index = static_cast*>(index.get()); + auto* spann_index = static_cast*>(index.get()); BOOST_REQUIRE(SSDServing::SSDIndex::Search(spann_index) == ErrorCode::Success); index.reset(); @@ -347,4 +357,225 @@ BOOST_AUTO_TEST_CASE(OfficialCompactRaBitQStoresRequestedBits) } } +BOOST_AUTO_TEST_CASE(SpannAppliesConfiguredADCWhenAttachingQuantizer) +{ + const auto raw = MakeRawVectors(); + auto quantizer = std::make_shared( + kDimension, kRaBitQBits, false); + BOOST_REQUIRE(quantizer->Train(raw) == ErrorCode::Success); + BOOST_CHECK(!quantizer->GetEnableADC()); + + auto index = VectorIndex::CreateInstance( + IndexAlgoType::SPANN, VectorValueType::UInt8); + BOOST_REQUIRE(index != nullptr); + index->SetParameter("EnableADC", "true", "BuildSSDIndex"); + index->SetQuantizer(quantizer); + BOOST_CHECK(quantizer->GetEnableADC()); +} + +BOOST_AUTO_TEST_CASE(RaBitQAutoTuneSelectsFirstQualifyingBit) +{ + std::vector evaluated; + int selected = 0; + float recall = 0.0F; + BOOST_REQUIRE( + COMMON::RaBitQAutoTuner::SelectMinimumBits( + 0.75F, + [&](int bits, float& value) { + evaluated.push_back(bits); + value = bits * 0.2F; + return ErrorCode::Success; + }, + selected, recall) == ErrorCode::Success); + BOOST_CHECK_EQUAL(selected, 4); + BOOST_CHECK_CLOSE(recall, 0.8F, 0.001F); + const std::vector expectedEvaluated = {4, 2, 3}; + BOOST_CHECK_EQUAL_COLLECTIONS( + evaluated.begin(), evaluated.end(), + expectedEvaluated.begin(), expectedEvaluated.end()); + + BOOST_CHECK( + COMMON::RaBitQAutoTuner::SelectMinimumBits( + 1.0F, + [](int, float& value) { + value = 0.99F; + return ErrorCode::Success; + }, + selected, recall) == ErrorCode::Fail); + BOOST_CHECK_EQUAL(selected, 0); +} + +BOOST_AUTO_TEST_CASE(RaBitQAutoTuneUsesDeeperTruthPool) +{ + const std::vector> truth = { + {10, 11, 12}, {20, 21, 22}}; + std::string error; + BOOST_CHECK( + COMMON::RaBitQAutoTuner::ValidateTruth( + truth, 32, 2, 2, error) == ErrorCode::Success); + BOOST_CHECK_EQUAL( + COMMON::RaBitQAutoTuner::RecallAtK(truth[0], {12, 10}, 2), 0.5F); + + BOOST_CHECK( + COMMON::RaBitQAutoTuner::ValidateTruth( + {{0, 1}, {2, 3}}, 4, 2, 2, error) != ErrorCode::Success); + BOOST_CHECK(error.find("greater than ResultNum") != std::string::npos); + BOOST_CHECK( + COMMON::RaBitQAutoTuner::ValidateTruth( + {{0, 1, 4}, {1, 2, 3}}, 4, 2, 2, error) != ErrorCode::Success); +} + +BOOST_AUTO_TEST_CASE(RaBitQStreamingCentroidUsesEveryVector) +{ + const auto raw = MakeRawVectors(); + auto oneShot = std::make_shared( + kDimension, kRaBitQBits, false); + BOOST_REQUIRE(oneShot->Train(raw) == ErrorCode::Success); + + auto streamed = std::make_shared( + kDimension, 1, false); + BOOST_REQUIRE(streamed->BeginTraining() == ErrorCode::Success); + const SizeType boundaries[] = {0, 7, 41, kVectorCount}; + for (std::size_t batch = 0; batch + 1 < std::size(boundaries); ++batch) { + ByteArray bytes = ByteArray::Alloc( + sizeof(float) * static_cast(boundaries[batch + 1] - boundaries[batch]) * + kDimension); + std::memcpy( + bytes.Data(), raw->GetVector(boundaries[batch]), + bytes.Length()); + auto batchVectors = std::make_shared( + bytes, VectorValueType::Float, kDimension, + boundaries[batch + 1] - boundaries[batch]); + BOOST_REQUIRE(streamed->AddTrainingBatch(batchVectors) == ErrorCode::Success); + } + BOOST_REQUIRE(streamed->FinishTraining() == ErrorCode::Success); + const auto sharedCentroid = streamed->CreateWithBits(kRaBitQBits); + BOOST_REQUIRE(sharedCentroid != nullptr); + + std::vector expected(oneShot->GetNumSubvectors()); + std::vector actual(sharedCentroid->GetNumSubvectors()); + oneShot->QuantizeVector(raw->GetVector(kVectorCount - 1), expected.data(), false); + sharedCentroid->QuantizeVector(raw->GetVector(kVectorCount - 1), actual.data(), false); + BOOST_CHECK_EQUAL_COLLECTIONS( + expected.begin(), expected.end(), actual.begin(), actual.end()); +} + +BOOST_AUTO_TEST_CASE(RaBitQEncodedWidthMatchesSavedModel) +{ + constexpr DimensionType dimension = 70; + constexpr int bits = 5; + ByteArray bytes = ByteArray::Alloc(sizeof(float) * dimension * 2); + auto* values = reinterpret_cast(bytes.Data()); + for (DimensionType i = 0; i < dimension * 2; ++i) { + values[i] = static_cast(i) / 13.0F; + } + auto vectors = std::make_shared( + bytes, VectorValueType::Float, dimension, 2); + auto quantizer = std::make_shared( + dimension, bits, false); + BOOST_REQUIRE(quantizer->Train(vectors) == ErrorCode::Success); + BOOST_CHECK_EQUAL( + quantizer->GetNumSubvectors(), 128 * bits / 8 + 5 * sizeof(float)); + + const char* modelPath = "rabitq_width_model.bin"; + auto output = f_createIO(); + BOOST_REQUIRE(output->Initialize(modelPath, std::ios::out | std::ios::binary)); + BOOST_REQUIRE(quantizer->SaveQuantizer(output) == ErrorCode::Success); + output->ShutDown(); + auto input = f_createIO(); + BOOST_REQUIRE(input->Initialize(modelPath, std::ios::in | std::ios::binary)); + auto loaded = COMMON::IQuantizer::LoadIQuantizer(input); + BOOST_REQUIRE(loaded != nullptr); + BOOST_CHECK_EQUAL(loaded->GetNumSubvectors(), quantizer->GetNumSubvectors()); + std::remove(modelPath); +} + +BOOST_AUTO_TEST_CASE(RaBitQAutoTuneProducesNativeBuildHandoff) +{ + constexpr SizeType vectorCount = 6; + constexpr SizeType queryCount = 2; + constexpr DimensionType dimension = 8; + const char* basePath = "rabitq_auto_base.bin"; + const char* queryPath = "rabitq_auto_queries.bin"; + const char* truthPath = "rabitq_auto_truth.bin"; + const char* outputFolder = "rabitq_auto_handoff"; + std::filesystem::remove_all(outputFolder); + + auto writeVectors = [](const char* path, SizeType count, DimensionType dim, float offset) { + std::ofstream output(path, std::ios::binary | std::ios::trunc); + output.write(reinterpret_cast(&count), sizeof(count)); + output.write(reinterpret_cast(&dim), sizeof(dim)); + for (SizeType row = 0; row < count; ++row) { + for (DimensionType column = 0; column < dim; ++column) { + const float value = offset + row * 0.5F + column * 0.01F; + output.write(reinterpret_cast(&value), sizeof(value)); + } + } + BOOST_REQUIRE(output.good()); + }; + writeVectors(basePath, vectorCount, dimension, 0.0F); + writeVectors(queryPath, queryCount, dimension, 0.2F); + { + std::ofstream truth(truthPath, std::ios::binary | std::ios::trunc); + const DimensionType depth = 3; + const std::int32_t truthQueryCount = queryCount; + truth.write( + reinterpret_cast(&truthQueryCount), sizeof(truthQueryCount)); + truth.write(reinterpret_cast(&depth), sizeof(depth)); + const std::int32_t ids[] = {0, 1, 2, 1, 0, 2}; + truth.write(reinterpret_cast(ids), sizeof(ids)); + const float distances[] = {0.0F, 1.0F, 2.0F, 0.0F, 1.0F, 2.0F}; + truth.write(reinterpret_cast(distances), sizeof(distances)); + BOOST_REQUIRE(truth.good()); + } + + Helper::IniReader config; + config.SetParameter("Base", "ValueType", "Float"); + config.SetParameter("Base", "DistCalcMethod", "L2"); + config.SetParameter("Base", "Dim", std::to_string(dimension)); + config.SetParameter("Base", "VectorPath", basePath); + config.SetParameter("Base", "VectorType", "DEFAULT"); + config.SetParameter("Base", "QueryPath", queryPath); + config.SetParameter("Base", "QueryType", "DEFAULT"); + config.SetParameter("Base", "TruthPath", truthPath); + config.SetParameter("Base", "TruthType", "DEFAULT"); + config.SetParameter("SearchSSDIndex", "QueryCountLimit", std::to_string(queryCount)); + config.SetParameter("SearchSSDIndex", "ResultNum", "1"); + config.SetParameter("BuildSSDIndex", "NumberOfThreads", "2"); + config.SetParameter("RaBitQAutoTune", "isExecute", "true"); + config.SetParameter("RaBitQAutoTune", "TargetRecall", "0"); + + COMMON::RaBitQAutoTuneResult result; + std::string error; + BOOST_REQUIRE_MESSAGE( + COMMON::RaBitQAutoTuner::Run( + config, outputFolder, result, error) == ErrorCode::Success, + error); + BOOST_CHECK_EQUAL(result.selectedBits, 1); + BOOST_CHECK_EQUAL(result.vectorCount, vectorCount); + BOOST_REQUIRE(result.quantizer != nullptr); + BOOST_CHECK_EQUAL( + result.codeDimension, result.quantizer->GetNumSubvectors()); + + auto readerOptions = std::make_shared( + VectorValueType::UInt8, result.codeDimension, VectorFileType::DEFAULT); + auto encodedReader = Helper::VectorSetReader::CreateInstance(readerOptions); + BOOST_REQUIRE(encodedReader->LoadFile(result.vectorPath) == ErrorCode::Success); + const auto encoded = encodedReader->GetVectorSet(); + BOOST_CHECK_EQUAL(encoded->Count(), vectorCount); + BOOST_CHECK_EQUAL(encoded->Dimension(), result.quantizer->GetNumSubvectors()); + + auto modelInput = f_createIO(); + BOOST_REQUIRE(modelInput->Initialize( + result.quantizerPath.c_str(), std::ios::in | std::ios::binary)); + auto loaded = COMMON::IQuantizer::LoadIQuantizer(modelInput); + BOOST_REQUIRE(loaded != nullptr); + BOOST_CHECK_EQUAL(loaded->GetNumSubvectors(), encoded->Dimension()); + + std::remove(basePath); + std::remove(queryPath); + std::remove(truthPath); + std::filesystem::remove_all(outputFolder); +} + BOOST_AUTO_TEST_SUITE_END() diff --git a/docs/RaBitQ_Global_Quantizer.md b/docs/RaBitQ_Global_Quantizer.md index fd09365b7..3d563cfac 100644 --- a/docs/RaBitQ_Global_Quantizer.md +++ b/docs/RaBitQ_Global_Quantizer.md @@ -34,5 +34,4 @@ multiple of 64. The current adapter supports the official L2 estimator; cosine distance is intentionally unsupported. `Script_AE/iniFile/build_SPANN_sift1m_rabitq3_global.ini` is the canonical -SIFT1M example. It uses STATIC postings containing the global RaBitQ codes; -keep `PostingQuantizer=None` because RaBitQ is already the global quantizer. +SIFT1M example. It uses STATIC postings containing the global RaBitQ codes.