diff --git a/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h b/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h index b4cec5fef..6a37fe48a 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h +++ b/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h @@ -57,6 +57,8 @@ INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_switchDeleteModes_Test) friend class BF16TieredTest; friend class FP16TieredTest; +friend class INT8TieredTest; +friend class CommonTypeMetricTieredTests_TestDataSizeTieredHNSW_Test; INDEX_TEST_FRIEND_CLASS(BM_VecSimBasics) INDEX_TEST_FRIEND_CLASS(BM_VecSimCommon) diff --git a/src/VecSim/index_factories/brute_force_factory.cpp b/src/VecSim/index_factories/brute_force_factory.cpp index e2919d75f..a7afb9c88 100644 --- a/src/VecSim/index_factories/brute_force_factory.cpp +++ b/src/VecSim/index_factories/brute_force_factory.cpp @@ -33,10 +33,12 @@ inline VecSimIndex *NewIndex_ChooseMultiOrSingle(const BFParams *params, static AbstractIndexInitParams NewAbstractInitParams(const VecSimParams *params) { const BFParams *bfParams = ¶ms->algoParams.bfParams; + size_t dataSize = VecSimParams_GetDataSize(bfParams->type, bfParams->dim, bfParams->metric); AbstractIndexInitParams abstractInitParams = {.allocator = VecSimAllocator::newVecsimAllocator(), .dim = bfParams->dim, .vecType = bfParams->type, + .dataSize = dataSize, .metric = bfParams->metric, .blockSize = bfParams->blockSize, .multi = bfParams->multi, @@ -52,32 +54,30 @@ VecSimIndex *NewIndex(const VecSimParams *params, bool is_normalized) { VecSimIndex *NewIndex(const BFParams *bfparams, const AbstractIndexInitParams &abstractInitParams, bool is_normalized) { - // If the index metric is Cosine, and is_normalized == true, we will skip normalizing vectors - // and query blobs. - VecSimMetric metric; - if (is_normalized && bfparams->metric == VecSimMetric_Cosine) { - metric = VecSimMetric_IP; - } else { - metric = bfparams->metric; - } + if (bfparams->type == VecSimType_FLOAT32) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, bfparams->dim); + abstractInitParams.allocator, bfparams->metric, bfparams->dim, is_normalized); return NewIndex_ChooseMultiOrSingle(bfparams, abstractInitParams, indexComponents); } else if (bfparams->type == VecSimType_FLOAT64) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, bfparams->dim); + abstractInitParams.allocator, bfparams->metric, bfparams->dim, is_normalized); return NewIndex_ChooseMultiOrSingle(bfparams, abstractInitParams, indexComponents); } else if (bfparams->type == VecSimType_BFLOAT16) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, bfparams->dim); + abstractInitParams.allocator, bfparams->metric, bfparams->dim, is_normalized); return NewIndex_ChooseMultiOrSingle(bfparams, abstractInitParams, indexComponents); } else if (bfparams->type == VecSimType_FLOAT16) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, bfparams->dim); + abstractInitParams.allocator, bfparams->metric, bfparams->dim, is_normalized); return NewIndex_ChooseMultiOrSingle(bfparams, abstractInitParams, indexComponents); + } else if (bfparams->type == VecSimType_INT8) { + IndexComponents indexComponents = CreateIndexComponents( + abstractInitParams.allocator, bfparams->metric, bfparams->dim, is_normalized); + return NewIndex_ChooseMultiOrSingle(bfparams, abstractInitParams, + indexComponents); } // If we got here something is wrong. @@ -117,6 +117,11 @@ size_t EstimateInitialSize(const BFParams *params, bool is_normalized) { } else if (params->type == VecSimType_FLOAT16) { est += EstimateComponentsMemory(params->metric, is_normalized); est += EstimateInitialSize_ChooseMultiOrSingle(params->multi); + } else if (params->type == VecSimType_INT8) { + est += EstimateComponentsMemory(params->metric, is_normalized); + est += EstimateInitialSize_ChooseMultiOrSingle(params->multi); + } else { + throw std::invalid_argument("Invalid params->type"); } est += sizeof(DataBlocksContainer) + allocations_overhead; diff --git a/src/VecSim/index_factories/components/components_factory.h b/src/VecSim/index_factories/components/components_factory.h index 6f4e984d1..13eb0eb3c 100644 --- a/src/VecSim/index_factories/components/components_factory.h +++ b/src/VecSim/index_factories/components/components_factory.h @@ -14,14 +14,24 @@ template IndexComponents -CreateIndexComponents(std::shared_ptr allocator, VecSimMetric metric, size_t dim) { +CreateIndexComponents(std::shared_ptr allocator, VecSimMetric metric, size_t dim, + bool is_normalized) { unsigned char alignment = 0; spaces::dist_func_t distFunc = spaces::GetDistFunc(metric, dim, &alignment); // Currently we have only one distance calculator implementation auto indexCalculator = new (allocator) DistanceCalculatorCommon(allocator, distFunc); - PreprocessorsContainerParams ppParams = {.metric = metric, .dim = dim, .alignment = alignment}; + // If the index metric is Cosine, and is_normalized == true, we will skip normalizing vectors + // and query blobs. + VecSimMetric pp_metric; + if (is_normalized && metric == VecSimMetric_Cosine) { + pp_metric = VecSimMetric_IP; + } else { + pp_metric = metric; + } + PreprocessorsContainerParams ppParams = { + .metric = pp_metric, .dim = dim, .alignment = alignment}; auto preprocessors = CreatePreprocessorsContainer(allocator, ppParams); return {indexCalculator, preprocessors}; diff --git a/src/VecSim/index_factories/hnsw_factory.cpp b/src/VecSim/index_factories/hnsw_factory.cpp index 58f7091f6..819899363 100644 --- a/src/VecSim/index_factories/hnsw_factory.cpp +++ b/src/VecSim/index_factories/hnsw_factory.cpp @@ -33,10 +33,14 @@ NewIndex_ChooseMultiOrSingle(const HNSWParams *params, static AbstractIndexInitParams NewAbstractInitParams(const VecSimParams *params) { const HNSWParams *hnswParams = ¶ms->algoParams.hnswParams; + + size_t dataSize = + VecSimParams_GetDataSize(hnswParams->type, hnswParams->dim, hnswParams->metric); AbstractIndexInitParams abstractInitParams = {.allocator = VecSimAllocator::newVecsimAllocator(), .dim = hnswParams->dim, .vecType = hnswParams->type, + .dataSize = dataSize, .metric = hnswParams->metric, .blockSize = hnswParams->blockSize, .multi = hnswParams->multi, @@ -48,36 +52,32 @@ VecSimIndex *NewIndex(const VecSimParams *params, bool is_normalized) { const HNSWParams *hnswParams = ¶ms->algoParams.hnswParams; AbstractIndexInitParams abstractInitParams = NewAbstractInitParams(params); - // If the index metric is Cosine, and is_normalized == true, we will skip normalizing vectors - // and query blobs. - VecSimMetric metric; - if (is_normalized && hnswParams->metric == VecSimMetric_Cosine) { - metric = VecSimMetric_IP; - } else { - metric = hnswParams->metric; - } - if (hnswParams->type == VecSimType_FLOAT32) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, hnswParams->dim); + abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized); return NewIndex_ChooseMultiOrSingle(hnswParams, abstractInitParams, indexComponents); } else if (hnswParams->type == VecSimType_FLOAT64) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, hnswParams->dim); + abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized); return NewIndex_ChooseMultiOrSingle(hnswParams, abstractInitParams, indexComponents); } else if (hnswParams->type == VecSimType_BFLOAT16) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, hnswParams->dim); + abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized); return NewIndex_ChooseMultiOrSingle(hnswParams, abstractInitParams, indexComponents); } else if (hnswParams->type == VecSimType_FLOAT16) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, hnswParams->dim); + abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized); return NewIndex_ChooseMultiOrSingle(hnswParams, abstractInitParams, indexComponents); + } else if (hnswParams->type == VecSimType_INT8) { + IndexComponents indexComponents = CreateIndexComponents( + abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized); + return NewIndex_ChooseMultiOrSingle(hnswParams, abstractInitParams, + indexComponents); } // If we got here something is wrong. @@ -114,6 +114,11 @@ size_t EstimateInitialSize(const HNSWParams *params, bool is_normalized) { } else if (params->type == VecSimType_FLOAT16) { est += EstimateComponentsMemory(params->metric, is_normalized); est += EstimateInitialSize_ChooseMultiOrSingle(params->multi); + } else if (params->type == VecSimType_INT8) { + est += EstimateComponentsMemory(params->metric, is_normalized); + est += EstimateInitialSize_ChooseMultiOrSingle(params->multi); + } else { + throw std::invalid_argument("Invalid params->type"); } return est; } @@ -203,34 +208,32 @@ VecSimIndex *NewIndex(const std::string &location, bool is_normalized) { VecSimParams vecsimParams = {.algo = VecSimAlgo_HNSWLIB, .algoParams = {.hnswParams = HNSWParams{params}}}; - VecSimMetric metric; - if (is_normalized && params.metric == VecSimMetric_Cosine) { - metric = VecSimMetric_IP; - } else { - metric = params.metric; - } - AbstractIndexInitParams abstractInitParams = NewAbstractInitParams(&vecsimParams); if (params.type == VecSimType_FLOAT32) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, abstractInitParams.dim); + abstractInitParams.allocator, params.metric, abstractInitParams.dim, is_normalized); return NewIndex_ChooseMultiOrSingle(input, ¶ms, abstractInitParams, indexComponents, version); } else if (params.type == VecSimType_FLOAT64) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, abstractInitParams.dim); + abstractInitParams.allocator, params.metric, abstractInitParams.dim, is_normalized); return NewIndex_ChooseMultiOrSingle(input, ¶ms, abstractInitParams, indexComponents, version); } else if (params.type == VecSimType_BFLOAT16) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, abstractInitParams.dim); + abstractInitParams.allocator, params.metric, abstractInitParams.dim, is_normalized); return NewIndex_ChooseMultiOrSingle(input, ¶ms, abstractInitParams, indexComponents, version); } else if (params.type == VecSimType_FLOAT16) { IndexComponents indexComponents = CreateIndexComponents( - abstractInitParams.allocator, metric, abstractInitParams.dim); + abstractInitParams.allocator, params.metric, abstractInitParams.dim, is_normalized); return NewIndex_ChooseMultiOrSingle(input, ¶ms, abstractInitParams, indexComponents, version); + } else if (params.type == VecSimType_INT8) { + IndexComponents indexComponents = CreateIndexComponents( + abstractInitParams.allocator, params.metric, abstractInitParams.dim, is_normalized); + return NewIndex_ChooseMultiOrSingle(input, ¶ms, abstractInitParams, + indexComponents, version); } else { auto bad_name = VecSimType_ToString(params.type); if (bad_name == nullptr) { diff --git a/src/VecSim/index_factories/tiered_factory.cpp b/src/VecSim/index_factories/tiered_factory.cpp index 68635e4ca..930630692 100644 --- a/src/VecSim/index_factories/tiered_factory.cpp +++ b/src/VecSim/index_factories/tiered_factory.cpp @@ -42,9 +42,12 @@ inline VecSimIndex *NewIndex(const TieredIndexParams *params) { BFParams bf_params = NewBFParams(params); std::shared_ptr flat_allocator = VecSimAllocator::newVecsimAllocator(); + size_t dataSize = VecSimParams_GetDataSize(bf_params.type, bf_params.dim, bf_params.metric); + AbstractIndexInitParams abstractInitParams = {.allocator = flat_allocator, .dim = bf_params.dim, .vecType = bf_params.type, + .dataSize = dataSize, .metric = bf_params.metric, .blockSize = bf_params.blockSize, .multi = bf_params.multi, @@ -80,6 +83,10 @@ inline size_t EstimateInitialSize(const TieredIndexParams *params) { est += sizeof(TieredHNSWIndex); } else if (hnsw_params.type == VecSimType_FLOAT16) { est += sizeof(TieredHNSWIndex); + } else if (hnsw_params.type == VecSimType_INT8) { + est += sizeof(TieredHNSWIndex); + } else { + throw std::invalid_argument("Invalid hnsw_params.type"); } return est; @@ -96,6 +103,8 @@ VecSimIndex *NewIndex(const TieredIndexParams *params) { return TieredHNSWFactory::NewIndex(params); } else if (type == VecSimType_FLOAT16) { return TieredHNSWFactory::NewIndex(params); + } else if (type == VecSimType_INT8) { + return TieredHNSWFactory::NewIndex(params); } return nullptr; // Invalid type. } diff --git a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_INT8.h b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_INT8.h index 7716d8ad7..35223d8b9 100644 --- a/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_INT8.h +++ b/src/VecSim/spaces/IP/IP_AVX512F_BW_VL_VNNI_INT8.h @@ -34,7 +34,7 @@ static inline int INT8_InnerProductImp(const void *pVect1v, const void *pVect2v, // Deal with remainder first. `dim` is more than 32, so we have at least one 32-int_8 block, // so mask loading is guaranteed to be safe if constexpr (residual % 32) { - __mmask32 mask = (1LU << (residual % 32)) - 1; + constexpr __mmask32 mask = (1LU << (residual % 32)) - 1; __m256i temp_a = _mm256_maskz_loadu_epi8(mask, pVect1); __m512i va = _mm512_cvtepi8_epi16(temp_a); pVect1 += residual % 32; diff --git a/src/VecSim/spaces/normalize/compute_norm.h b/src/VecSim/spaces/normalize/compute_norm.h new file mode 100644 index 000000000..d58139648 --- /dev/null +++ b/src/VecSim/spaces/normalize/compute_norm.h @@ -0,0 +1,25 @@ +/* + *Copyright Redis Ltd. 2021 - present + *Licensed under your choice of the Redis Source Available License 2.0 (RSALv2) or + *the Server Side Public License v1 (SSPLv1). + */ + +#pragma once + +#include + +namespace spaces { + +template +static inline float IntegralType_ComputeNorm(const DataType *vec, const size_t dim) { + int sum = 0; + + for (size_t i = 0; i < dim; i++) { + // No need to cast to int because c++ integer promotion ensures vec[i] is promoted to int + // before multiplication. + sum += vec[i] * vec[i]; + } + return sqrt(sum); +} + +} // namespace spaces diff --git a/src/VecSim/spaces/normalize/normalize_naive.h b/src/VecSim/spaces/normalize/normalize_naive.h index 119c19dcf..88967e39a 100644 --- a/src/VecSim/spaces/normalize/normalize_naive.h +++ b/src/VecSim/spaces/normalize/normalize_naive.h @@ -8,6 +8,7 @@ #include "VecSim/types/bfloat16.h" #include "VecSim/types/float16.h" +#include "compute_norm.h" #include #include @@ -73,4 +74,13 @@ static inline void float16_normalizeVector(void *vec, const size_t dim) { } } +static inline void int8_normalizeVector(void *vec, const size_t dim) { + int8_t *input_vector = static_cast(vec); + + float norm = IntegralType_ComputeNorm(input_vector, dim); + + // Store norm at the end of the vector. + *reinterpret_cast(input_vector + dim) = norm; +} + } // namespace spaces diff --git a/src/VecSim/spaces/spaces.cpp b/src/VecSim/spaces/spaces.cpp index 4385b5e94..c73ec997f 100644 --- a/src/VecSim/spaces/spaces.cpp +++ b/src/VecSim/spaces/spaces.cpp @@ -108,4 +108,10 @@ normalizeVector_f GetNormalizeFunc return float16_normalizeVector; } +/** The returned function computes the norm and stores it at the end of the given vector */ +template <> +normalizeVector_f GetNormalizeFunc(void) { + return int8_normalizeVector; +} + } // namespace spaces diff --git a/src/VecSim/utils/vec_utils.cpp b/src/VecSim/utils/vec_utils.cpp index 99160c247..cbe61338b 100644 --- a/src/VecSim/utils/vec_utils.cpp +++ b/src/VecSim/utils/vec_utils.cpp @@ -27,6 +27,7 @@ const char *VecSimCommonStrings::FLOAT32_STRING = "FLOAT32"; const char *VecSimCommonStrings::FLOAT64_STRING = "FLOAT64"; const char *VecSimCommonStrings::BFLOAT16_STRING = "BFLOAT16"; const char *VecSimCommonStrings::FLOAT16_STRING = "FLOAT16"; +const char *VecSimCommonStrings::INT8_STRING = "INT8"; const char *VecSimCommonStrings::INT32_STRING = "INT32"; const char *VecSimCommonStrings::INT64_STRING = "INT64"; @@ -147,6 +148,8 @@ const char *VecSimType_ToString(VecSimType vecsimType) { return VecSimCommonStrings::BFLOAT16_STRING; case VecSimType_FLOAT16: return VecSimCommonStrings::FLOAT16_STRING; + case VecSimType_INT8: + return VecSimCommonStrings::INT8_STRING; case VecSimType_INT32: return VecSimCommonStrings::INT32_STRING; case VecSimType_INT64: @@ -195,6 +198,8 @@ size_t VecSimType_sizeof(VecSimType type) { return sizeof(bfloat16); case VecSimType_FLOAT16: return sizeof(float16); + case VecSimType_INT8: + return sizeof(int8_t); case VecSimType_INT32: return sizeof(int32_t); case VecSimType_INT64: @@ -202,3 +207,11 @@ size_t VecSimType_sizeof(VecSimType type) { } return 0; } + +size_t VecSimParams_GetDataSize(VecSimType type, size_t dim, VecSimMetric metric) { + size_t dataSize = VecSimType_sizeof(type) * dim; + if (type == VecSimType_INT8 && metric == VecSimMetric_Cosine) { + dataSize += sizeof(float); // For the norm + } + return dataSize; +} diff --git a/src/VecSim/utils/vec_utils.h b/src/VecSim/utils/vec_utils.h index abb0c5688..18a5d1db3 100644 --- a/src/VecSim/utils/vec_utils.h +++ b/src/VecSim/utils/vec_utils.h @@ -27,6 +27,7 @@ struct VecSimCommonStrings { static const char *FLOAT64_STRING; static const char *BFLOAT16_STRING; static const char *FLOAT16_STRING; + static const char *INT8_STRING; static const char *INT32_STRING; static const char *INT64_STRING; @@ -90,3 +91,6 @@ const char *VecSimMetric_ToString(VecSimMetric vecsimMetric); const char *VecSimSearchMode_ToString(VecSearchMode vecsimSearchMode); size_t VecSimType_sizeof(VecSimType vecsimType); + +/** Returns the size in bytes of a stored or query vector */ +size_t VecSimParams_GetDataSize(VecSimType type, size_t dim, VecSimMetric metric); diff --git a/src/VecSim/vec_sim.cpp b/src/VecSim/vec_sim.cpp index 56912b07e..1a6d241fb 100644 --- a/src/VecSim/vec_sim.cpp +++ b/src/VecSim/vec_sim.cpp @@ -138,6 +138,9 @@ extern "C" void VecSim_Normalize(void *blob, size_t dim, VecSimType type) { spaces::GetNormalizeFunc()(blob, dim); } else if (type == VecSimType_FLOAT16) { spaces::GetNormalizeFunc()(blob, dim); + } else if (type == VecSimType_INT8) { + // assuming blob is large enough to store the norm at the end of the vector + spaces::GetNormalizeFunc()(blob, dim); } } diff --git a/src/VecSim/vec_sim_common.h b/src/VecSim/vec_sim_common.h index 943338dee..e8062484a 100644 --- a/src/VecSim/vec_sim_common.h +++ b/src/VecSim/vec_sim_common.h @@ -36,6 +36,7 @@ typedef enum { VecSimType_FLOAT64, VecSimType_BFLOAT16, VecSimType_FLOAT16, + VecSimType_INT8, VecSimType_INT32, VecSimType_INT64 } VecSimType; diff --git a/src/VecSim/vec_sim_debug.cpp b/src/VecSim/vec_sim_debug.cpp index 98cc05c91..395a3a9e0 100644 --- a/src/VecSim/vec_sim_debug.cpp +++ b/src/VecSim/vec_sim_debug.cpp @@ -32,6 +32,9 @@ extern "C" int VecSimDebug_GetElementNeighborsInHNSWGraph(VecSimIndex *index, si } else if (info.type == VecSimType_FLOAT16) { return dynamic_cast *>(index) ->getHNSWElementNeighbors(label, neighborsData); + } else if (info.type == VecSimType_INT8) { + return dynamic_cast *>(index)->getHNSWElementNeighbors( + label, neighborsData); } else { assert(false && "Invalid data type"); } @@ -48,6 +51,9 @@ extern "C" int VecSimDebug_GetElementNeighborsInHNSWGraph(VecSimIndex *index, si } else if (info.type == VecSimType_FLOAT16) { return dynamic_cast *>(index) ->getHNSWElementNeighbors(label, neighborsData); + } else if (info.type == VecSimType_INT8) { + return dynamic_cast *>(index)->getHNSWElementNeighbors( + label, neighborsData); } else { assert(false && "Invalid data type"); } diff --git a/src/VecSim/vec_sim_index.h b/src/VecSim/vec_sim_index.h index 08599951b..a5762c835 100644 --- a/src/VecSim/vec_sim_index.h +++ b/src/VecSim/vec_sim_index.h @@ -25,6 +25,7 @@ * @param allocator The allocator to use for the index. * @param dim The dimension of the vectors in the index. * @param vecType The type of the vectors in the index. + * @param dataSize The size of stored vectors in bytes. * @param metric The metric to use in the index. * @param blockSize The block size to use in the index. * @param multi Determines if the index should multi-index or not. @@ -34,6 +35,7 @@ struct AbstractIndexInitParams { std::shared_ptr allocator; size_t dim; VecSimType vecType; + size_t dataSize; VecSimMetric metric; size_t blockSize; bool multi; @@ -91,9 +93,6 @@ struct VecSimIndexAbstract : public VecSimIndexInterface { return info; } - spaces::normalizeVector_f - normalize_func; // A pointer to a normalization function of specific type. - public: /** * @brief Construct a new Vec Sim Index object @@ -102,12 +101,13 @@ struct VecSimIndexAbstract : public VecSimIndexInterface { VecSimIndexAbstract(const AbstractIndexInitParams ¶ms, const IndexComponents &components) : VecSimIndexInterface(params.allocator), dim(params.dim), vecType(params.vecType), - dataSize(dim * VecSimType_sizeof(vecType)), metric(params.metric), + dataSize(params.dataSize), metric(params.metric), blockSize(params.blockSize ? params.blockSize : DEFAULT_BLOCK_SIZE), indexCalculator(components.indexCalculator), preprocessors(components.preprocessors), alignment(preprocessors->getAlignment()), lastMode(EMPTY_MODE), isMulti(params.multi), - logCallbackCtx(params.logCtx), normalize_func(spaces::GetNormalizeFunc()) { + logCallbackCtx(params.logCtx) { assert(VecSimType_sizeof(vecType)); + assert(dataSize); } /** diff --git a/src/python_bindings/bindings.cpp b/src/python_bindings/bindings.cpp index 13215d6a4..b72670fb8 100644 --- a/src/python_bindings/bindings.cpp +++ b/src/python_bindings/bindings.cpp @@ -534,6 +534,7 @@ PYBIND11_MODULE(VecSim, m) { .value("VecSimType_FLOAT64", VecSimType_FLOAT64) .value("VecSimType_BFLOAT16", VecSimType_BFLOAT16) .value("VecSimType_FLOAT16", VecSimType_FLOAT16) + .value("VecSimType_INT8", VecSimType_INT8) .value("VecSimType_INT32", VecSimType_INT32) .value("VecSimType_INT64", VecSimType_INT64) .export_values(); diff --git a/tests/benchmark/spaces_benchmarks/bm_spaces_int8.cpp b/tests/benchmark/spaces_benchmarks/bm_spaces_int8.cpp index 0adde8972..25b7e85e0 100644 --- a/tests/benchmark/spaces_benchmarks/bm_spaces_int8.cpp +++ b/tests/benchmark/spaces_benchmarks/bm_spaces_int8.cpp @@ -28,8 +28,8 @@ class BM_VecSimSpaces_Integers_INT8 : public benchmark::Fixture { test_utils::populate_int8_vec(v2, dim, 1234); // Store the norm in the extra space for cosine calculations - *(float *)(v1 + dim) = test_utils::compute_norm(v1, dim); - *(float *)(v2 + dim) = test_utils::compute_norm(v2, dim); + *(float *)(v1 + dim) = test_utils::integral_compute_norm(v1, dim); + *(float *)(v2 + dim) = test_utils::integral_compute_norm(v2, dim); } void TearDown(const ::benchmark::State &state) { delete v1; diff --git a/tests/unit/CMakeLists.txt b/tests/unit/CMakeLists.txt index caa3fc522..5a4cc5108 100644 --- a/tests/unit/CMakeLists.txt +++ b/tests/unit/CMakeLists.txt @@ -31,14 +31,16 @@ endif() include(${root}/cmake/x86_64InstructionFlags.cmake) add_executable(test_hnsw ../utils/mock_thread_pool.cpp test_hnsw.cpp test_hnsw_multi.cpp test_hnsw_tiered.cpp unit_test_utils.cpp) -add_executable(test_hnsw_parallel test_hnsw_parallel.cpp unit_test_utils.cpp) -add_executable(test_bruteforce test_bruteforce.cpp test_bruteforce_multi.cpp unit_test_utils.cpp) -add_executable(test_allocator test_allocator.cpp unit_test_utils.cpp) +add_executable(test_hnsw_parallel test_hnsw_parallel.cpp ../utils/mock_thread_pool.cpp unit_test_utils.cpp) +add_executable(test_bruteforce test_bruteforce.cpp test_bruteforce_multi.cpp ../utils/mock_thread_pool.cpp unit_test_utils.cpp) +add_executable(test_allocator test_allocator.cpp ../utils/mock_thread_pool.cpp unit_test_utils.cpp) add_executable(test_spaces test_spaces.cpp) add_executable(test_types test_types.cpp) -add_executable(test_common ../utils/mock_thread_pool.cpp unit_test_utils.cpp test_common.cpp) +add_executable(test_common ../utils/mock_thread_pool.cpp test_common.cpp unit_test_utils.cpp) +add_executable(test_components test_components.cpp ../utils/mock_thread_pool.cpp unit_test_utils.cpp) add_executable(test_bf16 ../utils/mock_thread_pool.cpp test_bf16.cpp unit_test_utils.cpp) add_executable(test_fp16 ../utils/mock_thread_pool.cpp test_fp16.cpp unit_test_utils.cpp) +add_executable(test_int8 ../utils/mock_thread_pool.cpp test_int8.cpp unit_test_utils.cpp) target_link_libraries(test_hnsw PUBLIC gtest_main VectorSimilarity) target_link_libraries(test_hnsw_parallel PUBLIC gtest_main VectorSimilarity) @@ -46,9 +48,11 @@ target_link_libraries(test_bruteforce PUBLIC gtest_main VectorSimilarity) target_link_libraries(test_allocator PUBLIC gtest_main VectorSimilarity) target_link_libraries(test_spaces PUBLIC gtest_main VectorSimilarity) target_link_libraries(test_common PUBLIC gtest_main VectorSimilarity) +target_link_libraries(test_components PUBLIC gtest_main VectorSimilarity) target_link_libraries(test_types PUBLIC gtest_main VectorSimilarity) target_link_libraries(test_bf16 PUBLIC gtest_main VectorSimilarity) target_link_libraries(test_fp16 PUBLIC gtest_main VectorSimilarity) +target_link_libraries(test_int8 PUBLIC gtest_main VectorSimilarity) include(GoogleTest) @@ -58,6 +62,8 @@ gtest_discover_tests(test_bruteforce) gtest_discover_tests(test_allocator) gtest_discover_tests(test_spaces) gtest_discover_tests(test_common) +gtest_discover_tests(test_components) gtest_discover_tests(test_types) gtest_discover_tests(test_bf16 TEST_PREFIX BF16UNIT_) gtest_discover_tests(test_fp16 TEST_PREFIX FP16UNIT_) +gtest_discover_tests(test_int8 TEST_PREFIX INT8UNIT_) diff --git a/tests/unit/test_common.cpp b/tests/unit/test_common.cpp index bdfd6d9f2..e0ccd8d4c 100644 --- a/tests/unit/test_common.cpp +++ b/tests/unit/test_common.cpp @@ -15,6 +15,7 @@ #include "VecSim/algorithms/hnsw/hnsw.h" #include "VecSim/index_factories/hnsw_factory.h" #include "mock_thread_pool.h" +#include "tests_utils.h" #include "VecSim/index_factories/tiered_factory.h" #include "VecSim/spaces/spaces.h" #include "VecSim/types/bfloat16.h" @@ -625,579 +626,178 @@ TEST(CommonAPITest, NormalizeFloat16) { ASSERT_NEAR(1.0, norm, 0.001); } -class IndexCalculatorTest : public ::testing::Test {}; - -namespace dummyCalcultor { - -using DummyType = int; -using dummy_dist_func_t = DummyType (*)(int); +TEST(CommonAPITest, NormalizeInt8) { + size_t dim = 20; + int8_t v[dim + sizeof(float)]; -int dummyDistFunc(int value) { return value; } + test_utils::populate_int8_vec(v, dim); -template -class DistanceCalculatorDummy : public DistanceCalculatorInterface { -public: - DistanceCalculatorDummy(std::shared_ptr allocator, dummy_dist_func_t dist_func) - : DistanceCalculatorInterface(allocator, dist_func) {} + VecSim_Normalize(v, dim, VecSimType_INT8); - virtual DistType calcDistance(const void *v1, const void *v2, size_t dim) const { - return this->dist_func(7); + float res_norm = *(reinterpret_cast(v + dim)); + // Check that the normalized vector norm is 1. + float norm = 0; + for (size_t i = 0; i < dim; ++i) { + float val = v[i] / res_norm; + norm += val * val; } -}; - -} // namespace dummyCalcultor - -TEST(IndexCalculatorTest, TestIndexCalculator) { - - std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); - - // Test computer with a distance function signature different from dim(v1, v2, dim()). - using namespace dummyCalcultor; - auto distance_calculator = DistanceCalculatorDummy(allocator, dummyDistFunc); - ASSERT_EQ(distance_calculator.calcDistance(nullptr, nullptr, 0), 7); + ASSERT_FLOAT_EQ(norm, 1.0); } -class PreprocessorsTest : public ::testing::Test {}; - -namespace dummyPreprocessors { - -using DummyType = int; - -enum pp_mode { STORAGE_ONLY, QUERY_ONLY, BOTH, EMPTY }; - -// Dummy storage preprocessor -template -class DummyStoragePreprocessor : public PreprocessorInterface { -public: - DummyStoragePreprocessor(std::shared_ptr allocator, int value_to_add_storage, - int value_to_add_query = 0) - : PreprocessorInterface(allocator), value_to_add_storage(value_to_add_storage), - value_to_add_query(value_to_add_query) { - if (!value_to_add_query) - value_to_add_query = value_to_add_storage; - } - - void preprocess(const void *original_blob, void *&storage_blob, void *&query_blob, - size_t processed_bytes_count, unsigned char alignment) const override { - - this->preprocessForStorage(original_blob, storage_blob, processed_bytes_count); - } - - void preprocessForStorage(const void *original_blob, void *&blob, - size_t processed_bytes_count) const override { - // If the blob was not allocated yet, allocate it. - if (blob == nullptr) { - blob = this->allocator->allocate(processed_bytes_count); - memcpy(blob, original_blob, processed_bytes_count); - } - static_cast(blob)[0] += value_to_add_storage; - } - void preprocessQueryInPlace(void *blob, size_t processed_bytes_count, - unsigned char alignment) const override {} - void preprocessQuery(const void *original_blob, void *&blob, size_t processed_bytes_count, - unsigned char alignment) const override { - /* do nothing*/ - } - -private: - int value_to_add_storage; - int value_to_add_query; -}; - -// Dummy query preprocessor -template -class DummyQueryPreprocessor : public PreprocessorInterface { -public: - DummyQueryPreprocessor(std::shared_ptr allocator, int value_to_add_storage, - int _value_to_add_query = 0) - : PreprocessorInterface(allocator), value_to_add_storage(value_to_add_storage), - value_to_add_query(_value_to_add_query) { - if (!_value_to_add_query) - value_to_add_query = value_to_add_storage; - } - - void preprocess(const void *original_blob, void *&storage_blob, void *&query_blob, - size_t processed_bytes_count, unsigned char alignment) const override { - this->preprocessQuery(original_blob, query_blob, processed_bytes_count, alignment); - } - - void preprocessForStorage(const void *original_blob, void *&blob, - size_t processed_bytes_count) const override { - /* do nothing*/ - } - void preprocessQueryInPlace(void *blob, size_t processed_bytes_count, - unsigned char alignment) const override { - static_cast(blob)[0] += value_to_add_query; - } - void preprocessQuery(const void *original_blob, void *&blob, size_t processed_bytes_count, - unsigned char alignment) const override { - // If the blob was not allocated yet, allocate it. - if (blob == nullptr) { - blob = this->allocator->allocate_aligned(processed_bytes_count, alignment); - memcpy(blob, original_blob, processed_bytes_count); - } - static_cast(blob)[0] += value_to_add_query; - } - -private: - int value_to_add_storage; - int value_to_add_query; -}; - -// Dummy mixed preprocessor (precesses the blobs differently) -template -class DummyMixedPreprocessor : public PreprocessorInterface { -public: - DummyMixedPreprocessor(std::shared_ptr allocator, int value_to_add_storage, - int value_to_add_query) - : PreprocessorInterface(allocator), value_to_add_storage(value_to_add_storage), - value_to_add_query(value_to_add_query) {} - void preprocess(const void *original_blob, void *&storage_blob, void *&query_blob, - size_t processed_bytes_count, unsigned char alignment) const override { - - // One blob was already allocated by a previous preprocessor(s) that process both blobs the - // same. The blobs are pointing to the same memory, we need to allocate another memory slot - // to split them. - if ((storage_blob == query_blob) && (query_blob != nullptr)) { - storage_blob = this->allocator->allocate(processed_bytes_count); - memcpy(storage_blob, query_blob, processed_bytes_count); - } +class CommonTypeMetricTests : public testing::TestWithParam> { +protected: + template + void test_datasize(); - // Either both are nullptr or they are pointing to different memory slots. Both cases are - // handled by the designated functions. - this->preprocessForStorage(original_blob, storage_blob, processed_bytes_count); - this->preprocessQuery(original_blob, query_blob, processed_bytes_count, alignment); - } + template + void test_initial_size_estimation(); - void preprocessForStorage(const void *original_blob, void *&blob, - size_t processed_bytes_count) const override { - // If the blob was not allocated yet, allocate it. - if (blob == nullptr) { - blob = this->allocator->allocate(processed_bytes_count); - memcpy(blob, original_blob, processed_bytes_count); - } - static_cast(blob)[0] += value_to_add_storage; - } - void preprocessQueryInPlace(void *blob, size_t processed_bytes_count, - unsigned char alignment) const override {} - void preprocessQuery(const void *original_blob, void *&blob, size_t processed_bytes_count, - unsigned char alignment) const override { - // If the blob was not allocated yet, allocate it. - if (blob == nullptr) { - blob = this->allocator->allocate_aligned(processed_bytes_count, alignment); - memcpy(blob, original_blob, processed_bytes_count); - } - static_cast(blob)[0] += value_to_add_query; - } + virtual void TearDown() { VecSimIndex_Free(index); } -private: - int value_to_add_storage; - int value_to_add_query; + VecSimIndex *index; }; -} // namespace dummyPreprocessors - -TEST(PreprocessorsTest, PreprocessorsTestBasicAlignmentTest) { - using namespace dummyPreprocessors; - std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); - - unsigned char alignment = 5; - auto preprocessor = PreprocessorsContainerAbstract(allocator, alignment); - const int original_blob[4] = {1, 1, 1, 1}; - size_t processed_bytes_count = sizeof(original_blob); - { - auto aligned_query = preprocessor.preprocessQuery(original_blob, processed_bytes_count); - unsigned char address_alignment = (uintptr_t)(aligned_query.get()) % alignment; - ASSERT_EQ(address_alignment, 0); - } - - // The index computer is responsible for releasing the distance calculator. +template +void CommonTypeMetricTests::test_datasize() { + size_t dim = 4; + VecSimType type = std::get<0>(GetParam()); + VecSimMetric metric = std::get<1>(GetParam()); + algo_params params = {.dim = dim, .metric = metric}; + this->index = test_utils::CreateNewIndex(params, type); + size_t actual = test_utils::CalcVectorDataSize(index, type); + size_t expected = dim * VecSimType_sizeof(type); + if (type == VecSimType_INT8 && metric == VecSimMetric_Cosine) { + expected += sizeof(float); + } + ASSERT_EQ(actual, expected); } -template -void MultiPPContainerEmpty() { - using namespace dummyPreprocessors; - std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); - constexpr size_t dim = 4; - const int original_blob[dim] = {1, 2, 3, 4}; - const int original_blob_cpy[dim] = {1, 2, 3, 4}; - - constexpr size_t n_preprocessors = 3; - - auto multiPPContainer = - MultiPreprocessorsContainer(allocator, alignment); - - { - ProcessedBlobs processed_blobs = - multiPPContainer.preprocess(original_blob, sizeof(original_blob)); - // Original blob should not be changed - CompareVectors(original_blob, original_blob_cpy, dim); - - const void *storage_blob = processed_blobs.getStorageBlob(); - const void *query_blob = processed_blobs.getQueryBlob(); +TEST_P(CommonTypeMetricTests, TestDataSizeBF) { this->test_datasize(); } +TEST_P(CommonTypeMetricTests, TestDataSizeHNSW) { this->test_datasize(); } - // Storage blob should not be reallocated or changed - ASSERT_EQ(storage_blob, (const int *)original_blob); - CompareVectors(original_blob, (const int *)storage_blob, dim); +template +void CommonTypeMetricTests::test_initial_size_estimation() { + size_t dim = 4; + VecSimType type = std::get<0>(GetParam()); + VecSimMetric metric = std::get<1>(GetParam()); + algo_params params = {.dim = dim, .metric = metric}; + this->index = test_utils::CreateNewIndex(params, type); - // query blob *values* should not be changed - CompareVectors(original_blob, (const int *)query_blob, dim); + size_t estimation = EstimateInitialSize(params); + size_t actual = index->getAllocationSize(); - // If alignment is set the query blob address should be aligned to the specified alignment. - if constexpr (alignment) { - unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; - ASSERT_EQ(address_alignment, 0); - } - } + ASSERT_EQ(estimation, actual); } -TEST(PreprocessorsTest, MultiPPContainerEmptyNoAlignment) { - using namespace dummyPreprocessors; - MultiPPContainerEmpty<0>(); +TEST_P(CommonTypeMetricTests, TestInitialSizeEstimationBF) { + this->test_initial_size_estimation(); } - -TEST(PreprocessorsTest, MultiPPContainerEmptyAlignment) { - using namespace dummyPreprocessors; - MultiPPContainerEmpty<5>(); +TEST_P(CommonTypeMetricTests, TestInitialSizeEstimationHNSW) { + this->test_initial_size_estimation(); } -template -void MultiPreprocessorsContainerNoAlignment(dummyPreprocessors::pp_mode MODE) { - using namespace dummyPreprocessors; - std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); +class CommonTypeMetricTieredTests : public CommonTypeMetricTests { +protected: + virtual void TearDown() override {} - constexpr size_t n_preprocessors = 2; - unsigned char alignment = 0; - int initial_value = 1; - int value_to_add = 7; - const int original_blob[4] = {initial_value, initial_value, initial_value, initial_value}; - size_t processed_bytes_count = sizeof(original_blob); - - // Test computer with multiple preprocessors of the same type. - auto multiPPContainer = - MultiPreprocessorsContainer(allocator, alignment); - - auto verify_preprocess = [&](int expected_processed_value) { - ProcessedBlobs processed_blobs = - multiPPContainer.preprocess(original_blob, processed_bytes_count); - // Original blob should not be changed - ASSERT_EQ(original_blob[0], initial_value); - - const void *storage_blob = processed_blobs.getStorageBlob(); - const void *query_blob = processed_blobs.getQueryBlob(); - if (MODE == STORAGE_ONLY) { - // New storage blob should be allocated - ASSERT_NE(storage_blob, original_blob); - // query blob should be unprocessed - ASSERT_EQ(query_blob, original_blob); - ASSERT_EQ(((const int *)storage_blob)[0], expected_processed_value); - } else if (MODE == QUERY_ONLY) { - // New query blob should be allocated - ASSERT_NE(query_blob, original_blob); - // Storage blob should be unprocessed - ASSERT_EQ(storage_blob, original_blob); - ASSERT_EQ(((const int *)query_blob)[0], expected_processed_value); + tieredIndexMock mock_thread_pool; +}; + +TEST_P(CommonTypeMetricTieredTests, TestDataSizeTieredHNSW) { + size_t dim = 4; + VecSimType type = std::get<0>(GetParam()); + VecSimMetric metric = std::get<1>(GetParam()); + + HNSWParams hnsw_params = {.type = type, .dim = 4, .metric = metric}; + VecSimIndex *index = test_utils::CreateNewTieredHNSWIndex(hnsw_params, this->mock_thread_pool); + + auto verify_data_size = [&](const auto &tiered_index) { + auto hnsw_index = tiered_index->getHNSWIndex(); + auto bf_index = tiered_index->getFlatBufferIndex(); + size_t expected = dim * VecSimType_sizeof(type); + if (type == VecSimType_INT8 && metric == VecSimMetric_Cosine) { + expected += sizeof(float); } + size_t actual_hnsw = hnsw_index->getDataSize(); + ASSERT_EQ(actual_hnsw, expected); + size_t actual_bf = bf_index->getDataSize(); + ASSERT_EQ(actual_bf, expected); }; - /* ==== Add the first preprocessor ==== */ - auto preprocessor0 = new (allocator) PreprocessorType(allocator, value_to_add); - // add preprocessor returns next free spot in its preprocessors array. - ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor0), 1); - verify_preprocess(initial_value + value_to_add); - - /* ==== Add the second preprocessor ==== */ - auto preprocessor1 = new (allocator) PreprocessorType(allocator, value_to_add); - // add preprocessor returns 0 when adding the last preprocessor. - ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor1), 0); - ASSERT_NO_FATAL_FAILURE(verify_preprocess(initial_value + 2 * value_to_add)); -} - -TEST(PreprocessorsTest, MultiPreprocessorsContainerStorageNoAlignment) { - using namespace dummyPreprocessors; - MultiPreprocessorsContainerNoAlignment>( - pp_mode::STORAGE_ONLY); -} - -TEST(PreprocessorsTest, MultiPreprocessorsContainerQueryNoAlignment) { - using namespace dummyPreprocessors; - MultiPreprocessorsContainerNoAlignment>(pp_mode::QUERY_ONLY); -} - -template -void multiPPContainerMixedPreprocessorNoAlignment() { - using namespace dummyPreprocessors; - std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); - - constexpr size_t n_preprocessors = 3; - unsigned char alignment = 0; - int initial_value = 1; - int value_to_add_storage = 7; - int value_to_add_query = 2; - const int original_blob[4] = {initial_value, initial_value, initial_value, initial_value}; - size_t processed_bytes_count = sizeof(original_blob); - - // Test multiple preprocessors of the same type. - auto multiPPContainer = - MultiPreprocessorsContainer(allocator, alignment); - - /* ==== Add one preprocessor of each type ==== */ - auto preprocessor0 = - new (allocator) FirstPreprocessorType(allocator, value_to_add_storage, value_to_add_query); - ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor0), 1); - auto preprocessor1 = - new (allocator) SecondPreprocessorType(allocator, value_to_add_storage, value_to_add_query); - ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor1), 2); - - // scope this section so the blobs are released before the allocator. - { - ProcessedBlobs processed_blobs = - multiPPContainer.preprocess(original_blob, processed_bytes_count); - // Original blob should not be changed - ASSERT_EQ(original_blob[0], initial_value); - - // Both blobs should be allocated - const void *storage_blob = processed_blobs.getStorageBlob(); - const void *query_blob = processed_blobs.getQueryBlob(); - - // Ensure the computer process returns a new allocation of the expected processed blob with - // the new value. - ASSERT_NE(storage_blob, original_blob); - ASSERT_NE(query_blob, original_blob); - ASSERT_NE(query_blob, storage_blob); - - ASSERT_EQ(((const int *)storage_blob)[0], initial_value + value_to_add_storage); - ASSERT_EQ(((const int *)query_blob)[0], initial_value + value_to_add_query); + switch (type) { + case VecSimType_FLOAT32: { + auto tiered_index = test_utils::cast_to_tiered_index(index); + verify_data_size(tiered_index); + break; } - - /* ==== Add a preprocessor that processes both storage and query ==== */ - auto preprocessor2 = new (allocator) - DummyMixedPreprocessor(allocator, value_to_add_storage, value_to_add_query); - // add preprocessor returns 0 when adding the last preprocessor. - ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor2), 0); - { - ProcessedBlobs mixed_processed_blobs = - multiPPContainer.preprocess(original_blob, processed_bytes_count); - - const void *mixed_pp_storage_blob = mixed_processed_blobs.getStorageBlob(); - const void *mixed_pp_query_blob = mixed_processed_blobs.getQueryBlob(); - - // Ensure the computer process both blobs. - ASSERT_EQ(((const int *)mixed_pp_storage_blob)[0], - initial_value + 2 * value_to_add_storage); - ASSERT_EQ(((const int *)mixed_pp_query_blob)[0], initial_value + 2 * value_to_add_query); + case VecSimType_FLOAT64: { + auto tiered_index = test_utils::cast_to_tiered_index(index); + verify_data_size(tiered_index); + break; } - - // try adding another preprocessor and fail. - ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor2), -1); -} - -TEST(PreprocessorsTest, multiPPContainerMixedPreprocessorQueryFirst) { - using namespace dummyPreprocessors; - multiPPContainerMixedPreprocessorNoAlignment, - DummyStoragePreprocessor>(); -} - -TEST(PreprocessorsTest, multiPPContainerMixedPreprocessorStorageFirst) { - using namespace dummyPreprocessors; - multiPPContainerMixedPreprocessorNoAlignment, - DummyQueryPreprocessor>(); -} - -template -void multiPPContainerAlignment(dummyPreprocessors::pp_mode MODE) { - using namespace dummyPreprocessors; - std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); - - unsigned char alignment = 5; - constexpr size_t n_preprocessors = 1; - int initial_value = 1; - int value_to_add = 7; - const int original_blob[4] = {initial_value, initial_value, initial_value, initial_value}; - size_t processed_bytes_count = sizeof(original_blob); - - auto multiPPContainer = - MultiPreprocessorsContainer(allocator, alignment); - - auto verify_preprocess = [&](int expected_processed_value) { - ProcessedBlobs processed_blobs = - multiPPContainer.preprocess(original_blob, processed_bytes_count); - - const void *storage_blob = processed_blobs.getStorageBlob(); - const void *query_blob = processed_blobs.getQueryBlob(); - if (MODE == STORAGE_ONLY) { - // New storage blob should be allocated and processed - ASSERT_NE(storage_blob, original_blob); - ASSERT_EQ(((const int *)storage_blob)[0], expected_processed_value); - // query blob *values* should be unprocessed, however, it might be allocated if the - // original blob is not aligned. - ASSERT_EQ(((const int *)query_blob)[0], original_blob[0]); - } else if (MODE == QUERY_ONLY) { - // New query blob should be allocated - ASSERT_NE(query_blob, original_blob); - // Storage blob should be unprocessed and not allocated. - ASSERT_EQ(storage_blob, original_blob); - ASSERT_EQ(((const int *)query_blob)[0], expected_processed_value); - } - - // anyway the query blob should be aligned - unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; - ASSERT_EQ(address_alignment, 0); - }; - - auto preprocessor0 = new (allocator) PreprocessorType(allocator, value_to_add); - // add preprocessor returns next free spot in its preprocessors array. - ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor0), 0); - verify_preprocess(initial_value + value_to_add); -} - -TEST(PreprocessorsTest, StoragePreprocessorWithAlignment) { - using namespace dummyPreprocessors; - multiPPContainerAlignment>(pp_mode::STORAGE_ONLY); -} - -TEST(PreprocessorsTest, QueryPreprocessorWithAlignment) { - using namespace dummyPreprocessors; - multiPPContainerAlignment>(pp_mode::QUERY_ONLY); -} - -TEST(PreprocessorsTest, multiPPContainerCosineThenMixedPreprocess) { - using namespace dummyPreprocessors; - std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); - - constexpr size_t n_preprocessors = 2; - constexpr size_t dim = 4; - unsigned char alignment = 5; - - float initial_value = 1.0f; - float normalized_value = 0.5f; - float value_to_add_storage = 7.0f; - float value_to_add_query = 2.0f; - const float original_blob[dim] = {initial_value, initial_value, initial_value, initial_value}; - - auto multiPPContainer = - MultiPreprocessorsContainer(allocator, alignment); - - // adding cosine preprocessor - auto cosine_preprocessor = new (allocator) CosinePreprocessor(allocator, dim); - multiPPContainer.addPreprocessor(cosine_preprocessor); - { - ProcessedBlobs processed_blobs = - multiPPContainer.preprocess(original_blob, sizeof(original_blob)); - const void *storage_blob = processed_blobs.getStorageBlob(); - const void *query_blob = processed_blobs.getQueryBlob(); - // blobs should point to the same memory slot - ASSERT_EQ(storage_blob, query_blob); - // memory should be aligned - unsigned char address_alignment = (uintptr_t)(storage_blob) % alignment; - ASSERT_EQ(address_alignment, 0); - // They need to be allocated and processed - ASSERT_NE(storage_blob, nullptr); - ASSERT_EQ(((const float *)storage_blob)[0], normalized_value); - // the original blob should not change - ASSERT_NE(storage_blob, original_blob); + case VecSimType_BFLOAT16: { + auto tiered_index = test_utils::cast_to_tiered_index(index); + verify_data_size(tiered_index); + break; } - // adding mixed preprocessor - auto mixed_preprocessor = new (allocator) - DummyMixedPreprocessor(allocator, value_to_add_storage, value_to_add_query); - multiPPContainer.addPreprocessor(mixed_preprocessor); - { - ProcessedBlobs processed_blobs = - multiPPContainer.preprocess(original_blob, sizeof(original_blob)); - const void *storage_blob = processed_blobs.getStorageBlob(); - const void *query_blob = processed_blobs.getQueryBlob(); - // blobs should point to a different memory slot - ASSERT_NE(storage_blob, query_blob); - ASSERT_NE(storage_blob, nullptr); - ASSERT_NE(query_blob, nullptr); - - // query blob should be aligned - unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; - ASSERT_EQ(address_alignment, 0); - - // They need to be processed by both processors. - ASSERT_EQ(((const float *)storage_blob)[0], normalized_value + value_to_add_storage); - ASSERT_EQ(((const float *)query_blob)[0], normalized_value + value_to_add_query); - - // the original blob should not change - ASSERT_NE(storage_blob, original_blob); - ASSERT_NE(query_blob, original_blob); + case VecSimType_FLOAT16: { + auto tiered_index = test_utils::cast_to_tiered_index(index); + verify_data_size(tiered_index); + break; } - // The preprocessors should be released by the preprocessors container. -} - -TEST(PreprocessorsTest, multiPPContainerMixedThenCosinePreprocess) { - using namespace dummyPreprocessors; - std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); - - constexpr size_t n_preprocessors = 2; - constexpr size_t dim = 4; - unsigned char alignment = 5; - - float initial_value = 1.0f; - float normalized_value = 0.5f; - float value_to_add_storage = 7.0f; - float value_to_add_query = 2.0f; - const float original_blob[dim] = {initial_value, initial_value, initial_value, initial_value}; - - // Creating multi preprocessors container - auto mixed_preprocessor = new (allocator) - DummyMixedPreprocessor(allocator, value_to_add_storage, value_to_add_query); - auto multiPPContainer = - MultiPreprocessorsContainer(allocator, alignment); - multiPPContainer.addPreprocessor(mixed_preprocessor); - - { - ProcessedBlobs processed_blobs = - multiPPContainer.preprocess(original_blob, sizeof(original_blob)); - const void *storage_blob = processed_blobs.getStorageBlob(); - const void *query_blob = processed_blobs.getQueryBlob(); - // blobs should point to a different memory slot - ASSERT_NE(storage_blob, query_blob); - ASSERT_NE(storage_blob, nullptr); - ASSERT_NE(query_blob, nullptr); - - // query blob should be aligned - unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; - ASSERT_EQ(address_alignment, 0); - - // They need to be processed by both processors. - ASSERT_EQ(((const float *)storage_blob)[0], initial_value + value_to_add_storage); - ASSERT_EQ(((const float *)query_blob)[0], initial_value + value_to_add_query); - - // the original blob should not change - ASSERT_NE(storage_blob, original_blob); - ASSERT_NE(query_blob, original_blob); + case VecSimType_INT8: { + auto tiered_index = test_utils::cast_to_tiered_index(index); + verify_data_size(tiered_index); + break; } - - // adding cosine preprocessor - auto cosine_preprocessor = new (allocator) CosinePreprocessor(allocator, dim); - multiPPContainer.addPreprocessor(cosine_preprocessor); - { - ProcessedBlobs processed_blobs = - multiPPContainer.preprocess(original_blob, sizeof(original_blob)); - const void *storage_blob = processed_blobs.getStorageBlob(); - const void *query_blob = processed_blobs.getQueryBlob(); - // blobs should point to a different memory slot - ASSERT_NE(storage_blob, query_blob); - // query memory should be aligned - unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; - ASSERT_EQ(address_alignment, 0); - // They need to be allocated and processed - ASSERT_NE(storage_blob, nullptr); - ASSERT_NE(query_blob, nullptr); - float expected_processed_storage[dim] = {initial_value + value_to_add_storage, - initial_value, initial_value, initial_value}; - float expected_processed_query[dim] = {initial_value + value_to_add_query, initial_value, - initial_value, initial_value}; - VecSim_Normalize(expected_processed_storage, dim, VecSimType_FLOAT32); - VecSim_Normalize(expected_processed_query, dim, VecSimType_FLOAT32); - ASSERT_EQ(((const float *)storage_blob)[0], expected_processed_storage[0]); - ASSERT_EQ(((const float *)query_blob)[0], expected_processed_query[0]); - // the original blob should not change - ASSERT_NE(storage_blob, original_blob); - ASSERT_NE(query_blob, original_blob); + default: + FAIL() << "Unsupported data type"; } - // The preprocessors should be released by the preprocessors container. } + +TEST_P(CommonTypeMetricTieredTests, TestInitialSizeEstimationTieredHNSW) { + size_t dim = 4; + VecSimType type = std::get<0>(GetParam()); + VecSimMetric metric = std::get<1>(GetParam()); + HNSWParams hnsw_params = {.type = type, .dim = dim, .metric = metric}; + VecSimParams vecsim_hnsw_params = CreateParams(hnsw_params); + TieredIndexParams tiered_params = + test_utils::CreateTieredParams(vecsim_hnsw_params, this->mock_thread_pool); + VecSimParams params = CreateParams(tiered_params); + auto *index = VecSimIndex_New(¶ms); + mock_thread_pool.ctx->index_strong_ref.reset(index); + + size_t estimation = VecSimIndex_EstimateInitialSize(¶ms); + size_t actual = index->getAllocationSize(); + + ASSERT_EQ(estimation, actual); +} + +constexpr VecSimType vecsim_datatypes[] = {VecSimType_FLOAT32, VecSimType_FLOAT64, + VecSimType_BFLOAT16, VecSimType_FLOAT16, + VecSimType_INT8}; + +/** Run all CommonTypeMetricTests tests for each {VecSimType, VecSimMetric} combination */ +INSTANTIATE_TEST_SUITE_P(CommonTest, CommonTypeMetricTests, + testing::Combine(testing::ValuesIn(vecsim_datatypes), + testing::Values(VecSimMetric_L2, VecSimMetric_IP, + VecSimMetric_Cosine)), + [](const testing::TestParamInfo &info) { + const char *type = VecSimType_ToString(std::get<0>(info.param)); + const char *metric = VecSimMetric_ToString(std::get<1>(info.param)); + std::string test_name(type); + return test_name + "_" + metric; + }); + +/** Run all CommonTypeMetricTieredTests tests for each {VecSimType, VecSimMetric} combination */ +INSTANTIATE_TEST_SUITE_P( + CommonTieredTest, CommonTypeMetricTieredTests, + testing::Combine(testing::ValuesIn(vecsim_datatypes), + testing::Values(VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine)), + [](const testing::TestParamInfo &info) { + const char *type = VecSimType_ToString(std::get<0>(info.param)); + const char *metric = VecSimMetric_ToString(std::get<1>(info.param)); + std::string test_name(type); + return test_name + "_" + metric; + }); diff --git a/tests/unit/test_components.cpp b/tests/unit/test_components.cpp new file mode 100644 index 000000000..af49b12a8 --- /dev/null +++ b/tests/unit/test_components.cpp @@ -0,0 +1,587 @@ +/* + *Copyright Redis Ltd. 2021 - present + *Licensed under your choice of the Redis Source Available License 2.0 (RSALv2) or + *the Server Side Public License v1 (SSPLv1). + */ + +#include "gtest/gtest.h" +#include "VecSim/vec_sim.h" +#include "VecSim/spaces/computer/preprocessor_container.h" +#include "VecSim/spaces/computer/calculator.h" +#include "unit_test_utils.h" + +class IndexCalculatorTest : public ::testing::Test {}; +namespace dummyCalcultor { + +using DummyType = int; +using dummy_dist_func_t = DummyType (*)(int); + +int dummyDistFunc(int value) { return value; } + +template +class DistanceCalculatorDummy : public DistanceCalculatorInterface { +public: + DistanceCalculatorDummy(std::shared_ptr allocator, dummy_dist_func_t dist_func) + : DistanceCalculatorInterface(allocator, dist_func) {} + + virtual DistType calcDistance(const void *v1, const void *v2, size_t dim) const { + return this->dist_func(7); + } +}; + +} // namespace dummyCalcultor + +TEST(IndexCalculatorTest, TestIndexCalculator) { + + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + + // Test computer with a distance function signature different from dim(v1, v2, dim()). + using namespace dummyCalcultor; + auto distance_calculator = DistanceCalculatorDummy(allocator, dummyDistFunc); + + ASSERT_EQ(distance_calculator.calcDistance(nullptr, nullptr, 0), 7); +} + +class PreprocessorsTest : public ::testing::Test {}; + +namespace dummyPreprocessors { + +using DummyType = int; + +enum pp_mode { STORAGE_ONLY, QUERY_ONLY, BOTH, EMPTY }; + +// Dummy storage preprocessor +template +class DummyStoragePreprocessor : public PreprocessorInterface { +public: + DummyStoragePreprocessor(std::shared_ptr allocator, int value_to_add_storage, + int value_to_add_query = 0) + : PreprocessorInterface(allocator), value_to_add_storage(value_to_add_storage), + value_to_add_query(value_to_add_query) { + if (!value_to_add_query) + value_to_add_query = value_to_add_storage; + } + + void preprocess(const void *original_blob, void *&storage_blob, void *&query_blob, + size_t processed_bytes_count, unsigned char alignment) const override { + + this->preprocessForStorage(original_blob, storage_blob, processed_bytes_count); + } + + void preprocessForStorage(const void *original_blob, void *&blob, + size_t processed_bytes_count) const override { + // If the blob was not allocated yet, allocate it. + if (blob == nullptr) { + blob = this->allocator->allocate(processed_bytes_count); + memcpy(blob, original_blob, processed_bytes_count); + } + static_cast(blob)[0] += value_to_add_storage; + } + void preprocessQueryInPlace(void *blob, size_t processed_bytes_count, + unsigned char alignment) const override {} + void preprocessQuery(const void *original_blob, void *&blob, size_t processed_bytes_count, + unsigned char alignment) const override { + /* do nothing*/ + } + +private: + int value_to_add_storage; + int value_to_add_query; +}; + +// Dummy query preprocessor +template +class DummyQueryPreprocessor : public PreprocessorInterface { +public: + DummyQueryPreprocessor(std::shared_ptr allocator, int value_to_add_storage, + int _value_to_add_query = 0) + : PreprocessorInterface(allocator), value_to_add_storage(value_to_add_storage), + value_to_add_query(_value_to_add_query) { + if (!_value_to_add_query) + value_to_add_query = value_to_add_storage; + } + + void preprocess(const void *original_blob, void *&storage_blob, void *&query_blob, + size_t processed_bytes_count, unsigned char alignment) const override { + this->preprocessQuery(original_blob, query_blob, processed_bytes_count, alignment); + } + + void preprocessForStorage(const void *original_blob, void *&blob, + size_t processed_bytes_count) const override { + /* do nothing*/ + } + void preprocessQueryInPlace(void *blob, size_t processed_bytes_count, + unsigned char alignment) const override { + static_cast(blob)[0] += value_to_add_query; + } + void preprocessQuery(const void *original_blob, void *&blob, size_t processed_bytes_count, + unsigned char alignment) const override { + // If the blob was not allocated yet, allocate it. + if (blob == nullptr) { + blob = this->allocator->allocate_aligned(processed_bytes_count, alignment); + memcpy(blob, original_blob, processed_bytes_count); + } + static_cast(blob)[0] += value_to_add_query; + } + +private: + int value_to_add_storage; + int value_to_add_query; +}; + +// Dummy mixed preprocessor (precesses the blobs differently) +template +class DummyMixedPreprocessor : public PreprocessorInterface { +public: + DummyMixedPreprocessor(std::shared_ptr allocator, int value_to_add_storage, + int value_to_add_query) + : PreprocessorInterface(allocator), value_to_add_storage(value_to_add_storage), + value_to_add_query(value_to_add_query) {} + void preprocess(const void *original_blob, void *&storage_blob, void *&query_blob, + size_t processed_bytes_count, unsigned char alignment) const override { + + // One blob was already allocated by a previous preprocessor(s) that process both blobs the + // same. The blobs are pointing to the same memory, we need to allocate another memory slot + // to split them. + if ((storage_blob == query_blob) && (query_blob != nullptr)) { + storage_blob = this->allocator->allocate(processed_bytes_count); + memcpy(storage_blob, query_blob, processed_bytes_count); + } + + // Either both are nullptr or they are pointing to different memory slots. Both cases are + // handled by the designated functions. + this->preprocessForStorage(original_blob, storage_blob, processed_bytes_count); + this->preprocessQuery(original_blob, query_blob, processed_bytes_count, alignment); + } + + void preprocessForStorage(const void *original_blob, void *&blob, + size_t processed_bytes_count) const override { + // If the blob was not allocated yet, allocate it. + if (blob == nullptr) { + blob = this->allocator->allocate(processed_bytes_count); + memcpy(blob, original_blob, processed_bytes_count); + } + static_cast(blob)[0] += value_to_add_storage; + } + void preprocessQueryInPlace(void *blob, size_t processed_bytes_count, + unsigned char alignment) const override {} + void preprocessQuery(const void *original_blob, void *&blob, size_t processed_bytes_count, + unsigned char alignment) const override { + // If the blob was not allocated yet, allocate it. + if (blob == nullptr) { + blob = this->allocator->allocate_aligned(processed_bytes_count, alignment); + memcpy(blob, original_blob, processed_bytes_count); + } + static_cast(blob)[0] += value_to_add_query; + } + +private: + int value_to_add_storage; + int value_to_add_query; +}; +} // namespace dummyPreprocessors + +TEST(PreprocessorsTest, PreprocessorsTestBasicAlignmentTest) { + using namespace dummyPreprocessors; + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + + unsigned char alignment = 5; + auto preprocessor = PreprocessorsContainerAbstract(allocator, alignment); + const int original_blob[4] = {1, 1, 1, 1}; + size_t processed_bytes_count = sizeof(original_blob); + + { + auto aligned_query = preprocessor.preprocessQuery(original_blob, processed_bytes_count); + unsigned char address_alignment = (uintptr_t)(aligned_query.get()) % alignment; + ASSERT_EQ(address_alignment, 0); + } + + // The index computer is responsible for releasing the distance calculator. +} + +template +void MultiPPContainerEmpty() { + using namespace dummyPreprocessors; + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 4; + const int original_blob[dim] = {1, 2, 3, 4}; + const int original_blob_cpy[dim] = {1, 2, 3, 4}; + + constexpr size_t n_preprocessors = 3; + + auto multiPPContainer = + MultiPreprocessorsContainer(allocator, alignment); + + { + ProcessedBlobs processed_blobs = + multiPPContainer.preprocess(original_blob, sizeof(original_blob)); + // Original blob should not be changed + CompareVectors(original_blob, original_blob_cpy, dim); + + const void *storage_blob = processed_blobs.getStorageBlob(); + const void *query_blob = processed_blobs.getQueryBlob(); + + // Storage blob should not be reallocated or changed + ASSERT_EQ(storage_blob, (const int *)original_blob); + CompareVectors(original_blob, (const int *)storage_blob, dim); + + // query blob *values* should not be changed + CompareVectors(original_blob, (const int *)query_blob, dim); + + // If alignment is set the query blob address should be aligned to the specified alignment. + if constexpr (alignment) { + unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; + ASSERT_EQ(address_alignment, 0); + } + } +} + +TEST(PreprocessorsTest, MultiPPContainerEmptyNoAlignment) { + using namespace dummyPreprocessors; + MultiPPContainerEmpty<0>(); +} + +TEST(PreprocessorsTest, MultiPPContainerEmptyAlignment) { + using namespace dummyPreprocessors; + MultiPPContainerEmpty<5>(); +} + +template +void MultiPreprocessorsContainerNoAlignment(dummyPreprocessors::pp_mode MODE) { + using namespace dummyPreprocessors; + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + + constexpr size_t n_preprocessors = 2; + unsigned char alignment = 0; + int initial_value = 1; + int value_to_add = 7; + const int original_blob[4] = {initial_value, initial_value, initial_value, initial_value}; + size_t processed_bytes_count = sizeof(original_blob); + + // Test computer with multiple preprocessors of the same type. + auto multiPPContainer = + MultiPreprocessorsContainer(allocator, alignment); + + auto verify_preprocess = [&](int expected_processed_value) { + ProcessedBlobs processed_blobs = + multiPPContainer.preprocess(original_blob, processed_bytes_count); + // Original blob should not be changed + ASSERT_EQ(original_blob[0], initial_value); + + const void *storage_blob = processed_blobs.getStorageBlob(); + const void *query_blob = processed_blobs.getQueryBlob(); + if (MODE == STORAGE_ONLY) { + // New storage blob should be allocated + ASSERT_NE(storage_blob, original_blob); + // query blob should be unprocessed + ASSERT_EQ(query_blob, original_blob); + ASSERT_EQ(((const int *)storage_blob)[0], expected_processed_value); + } else if (MODE == QUERY_ONLY) { + // New query blob should be allocated + ASSERT_NE(query_blob, original_blob); + // Storage blob should be unprocessed + ASSERT_EQ(storage_blob, original_blob); + ASSERT_EQ(((const int *)query_blob)[0], expected_processed_value); + } + }; + + /* ==== Add the first preprocessor ==== */ + auto preprocessor0 = new (allocator) PreprocessorType(allocator, value_to_add); + // add preprocessor returns next free spot in its preprocessors array. + ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor0), 1); + verify_preprocess(initial_value + value_to_add); + + /* ==== Add the second preprocessor ==== */ + auto preprocessor1 = new (allocator) PreprocessorType(allocator, value_to_add); + // add preprocessor returns 0 when adding the last preprocessor. + ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor1), 0); + ASSERT_NO_FATAL_FAILURE(verify_preprocess(initial_value + 2 * value_to_add)); +} + +TEST(PreprocessorsTest, MultiPreprocessorsContainerStorageNoAlignment) { + using namespace dummyPreprocessors; + MultiPreprocessorsContainerNoAlignment>( + pp_mode::STORAGE_ONLY); +} + +TEST(PreprocessorsTest, MultiPreprocessorsContainerQueryNoAlignment) { + using namespace dummyPreprocessors; + MultiPreprocessorsContainerNoAlignment>(pp_mode::QUERY_ONLY); +} + +template +void multiPPContainerMixedPreprocessorNoAlignment() { + using namespace dummyPreprocessors; + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + + constexpr size_t n_preprocessors = 3; + unsigned char alignment = 0; + int initial_value = 1; + int value_to_add_storage = 7; + int value_to_add_query = 2; + const int original_blob[4] = {initial_value, initial_value, initial_value, initial_value}; + size_t processed_bytes_count = sizeof(original_blob); + + // Test multiple preprocessors of the same type. + auto multiPPContainer = + MultiPreprocessorsContainer(allocator, alignment); + + /* ==== Add one preprocessor of each type ==== */ + auto preprocessor0 = + new (allocator) FirstPreprocessorType(allocator, value_to_add_storage, value_to_add_query); + ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor0), 1); + auto preprocessor1 = + new (allocator) SecondPreprocessorType(allocator, value_to_add_storage, value_to_add_query); + ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor1), 2); + + // scope this section so the blobs are released before the allocator. + { + ProcessedBlobs processed_blobs = + multiPPContainer.preprocess(original_blob, processed_bytes_count); + // Original blob should not be changed + ASSERT_EQ(original_blob[0], initial_value); + + // Both blobs should be allocated + const void *storage_blob = processed_blobs.getStorageBlob(); + const void *query_blob = processed_blobs.getQueryBlob(); + + // Ensure the computer process returns a new allocation of the expected processed blob with + // the new value. + ASSERT_NE(storage_blob, original_blob); + ASSERT_NE(query_blob, original_blob); + ASSERT_NE(query_blob, storage_blob); + + ASSERT_EQ(((const int *)storage_blob)[0], initial_value + value_to_add_storage); + ASSERT_EQ(((const int *)query_blob)[0], initial_value + value_to_add_query); + } + + /* ==== Add a preprocessor that processes both storage and query ==== */ + auto preprocessor2 = new (allocator) + DummyMixedPreprocessor(allocator, value_to_add_storage, value_to_add_query); + // add preprocessor returns 0 when adding the last preprocessor. + ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor2), 0); + { + ProcessedBlobs mixed_processed_blobs = + multiPPContainer.preprocess(original_blob, processed_bytes_count); + + const void *mixed_pp_storage_blob = mixed_processed_blobs.getStorageBlob(); + const void *mixed_pp_query_blob = mixed_processed_blobs.getQueryBlob(); + + // Ensure the computer process both blobs. + ASSERT_EQ(((const int *)mixed_pp_storage_blob)[0], + initial_value + 2 * value_to_add_storage); + ASSERT_EQ(((const int *)mixed_pp_query_blob)[0], initial_value + 2 * value_to_add_query); + } + + // try adding another preprocessor and fail. + ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor2), -1); +} + +TEST(PreprocessorsTest, multiPPContainerMixedPreprocessorQueryFirst) { + using namespace dummyPreprocessors; + multiPPContainerMixedPreprocessorNoAlignment, + DummyStoragePreprocessor>(); +} + +TEST(PreprocessorsTest, multiPPContainerMixedPreprocessorStorageFirst) { + using namespace dummyPreprocessors; + multiPPContainerMixedPreprocessorNoAlignment, + DummyQueryPreprocessor>(); +} + +template +void multiPPContainerAlignment(dummyPreprocessors::pp_mode MODE) { + using namespace dummyPreprocessors; + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + + unsigned char alignment = 5; + constexpr size_t n_preprocessors = 1; + int initial_value = 1; + int value_to_add = 7; + const int original_blob[4] = {initial_value, initial_value, initial_value, initial_value}; + size_t processed_bytes_count = sizeof(original_blob); + + auto multiPPContainer = + MultiPreprocessorsContainer(allocator, alignment); + + auto verify_preprocess = [&](int expected_processed_value) { + ProcessedBlobs processed_blobs = + multiPPContainer.preprocess(original_blob, processed_bytes_count); + + const void *storage_blob = processed_blobs.getStorageBlob(); + const void *query_blob = processed_blobs.getQueryBlob(); + if (MODE == STORAGE_ONLY) { + // New storage blob should be allocated and processed + ASSERT_NE(storage_blob, original_blob); + ASSERT_EQ(((const int *)storage_blob)[0], expected_processed_value); + // query blob *values* should be unprocessed, however, it might be allocated if the + // original blob is not aligned. + ASSERT_EQ(((const int *)query_blob)[0], original_blob[0]); + } else if (MODE == QUERY_ONLY) { + // New query blob should be allocated + ASSERT_NE(query_blob, original_blob); + // Storage blob should be unprocessed and not allocated. + ASSERT_EQ(storage_blob, original_blob); + ASSERT_EQ(((const int *)query_blob)[0], expected_processed_value); + } + + // anyway the query blob should be aligned + unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; + ASSERT_EQ(address_alignment, 0); + }; + + auto preprocessor0 = new (allocator) PreprocessorType(allocator, value_to_add); + // add preprocessor returns next free spot in its preprocessors array. + ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor0), 0); + verify_preprocess(initial_value + value_to_add); +} + +TEST(PreprocessorsTest, StoragePreprocessorWithAlignment) { + using namespace dummyPreprocessors; + multiPPContainerAlignment>(pp_mode::STORAGE_ONLY); +} + +TEST(PreprocessorsTest, QueryPreprocessorWithAlignment) { + using namespace dummyPreprocessors; + multiPPContainerAlignment>(pp_mode::QUERY_ONLY); +} + +TEST(PreprocessorsTest, multiPPContainerCosineThenMixedPreprocess) { + using namespace dummyPreprocessors; + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + + constexpr size_t n_preprocessors = 2; + constexpr size_t dim = 4; + unsigned char alignment = 5; + + float initial_value = 1.0f; + float normalized_value = 0.5f; + float value_to_add_storage = 7.0f; + float value_to_add_query = 2.0f; + const float original_blob[dim] = {initial_value, initial_value, initial_value, initial_value}; + + auto multiPPContainer = + MultiPreprocessorsContainer(allocator, alignment); + + // adding cosine preprocessor + auto cosine_preprocessor = new (allocator) CosinePreprocessor(allocator, dim); + multiPPContainer.addPreprocessor(cosine_preprocessor); + { + ProcessedBlobs processed_blobs = + multiPPContainer.preprocess(original_blob, sizeof(original_blob)); + const void *storage_blob = processed_blobs.getStorageBlob(); + const void *query_blob = processed_blobs.getQueryBlob(); + // blobs should point to the same memory slot + ASSERT_EQ(storage_blob, query_blob); + // memory should be aligned + unsigned char address_alignment = (uintptr_t)(storage_blob) % alignment; + ASSERT_EQ(address_alignment, 0); + // They need to be allocated and processed + ASSERT_NE(storage_blob, nullptr); + ASSERT_EQ(((const float *)storage_blob)[0], normalized_value); + // the original blob should not change + ASSERT_NE(storage_blob, original_blob); + } + // adding mixed preprocessor + auto mixed_preprocessor = new (allocator) + DummyMixedPreprocessor(allocator, value_to_add_storage, value_to_add_query); + multiPPContainer.addPreprocessor(mixed_preprocessor); + { + ProcessedBlobs processed_blobs = + multiPPContainer.preprocess(original_blob, sizeof(original_blob)); + const void *storage_blob = processed_blobs.getStorageBlob(); + const void *query_blob = processed_blobs.getQueryBlob(); + // blobs should point to a different memory slot + ASSERT_NE(storage_blob, query_blob); + ASSERT_NE(storage_blob, nullptr); + ASSERT_NE(query_blob, nullptr); + + // query blob should be aligned + unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; + ASSERT_EQ(address_alignment, 0); + + // They need to be processed by both processors. + ASSERT_EQ(((const float *)storage_blob)[0], normalized_value + value_to_add_storage); + ASSERT_EQ(((const float *)query_blob)[0], normalized_value + value_to_add_query); + + // the original blob should not change + ASSERT_NE(storage_blob, original_blob); + ASSERT_NE(query_blob, original_blob); + } + // The preprocessors should be released by the preprocessors container. +} + +TEST(PreprocessorsTest, multiPPContainerMixedThenCosinePreprocess) { + using namespace dummyPreprocessors; + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + + constexpr size_t n_preprocessors = 2; + constexpr size_t dim = 4; + unsigned char alignment = 5; + + float initial_value = 1.0f; + float normalized_value = 0.5f; + float value_to_add_storage = 7.0f; + float value_to_add_query = 2.0f; + const float original_blob[dim] = {initial_value, initial_value, initial_value, initial_value}; + + // Creating multi preprocessors container + auto mixed_preprocessor = new (allocator) + DummyMixedPreprocessor(allocator, value_to_add_storage, value_to_add_query); + auto multiPPContainer = + MultiPreprocessorsContainer(allocator, alignment); + multiPPContainer.addPreprocessor(mixed_preprocessor); + + { + ProcessedBlobs processed_blobs = + multiPPContainer.preprocess(original_blob, sizeof(original_blob)); + const void *storage_blob = processed_blobs.getStorageBlob(); + const void *query_blob = processed_blobs.getQueryBlob(); + // blobs should point to a different memory slot + ASSERT_NE(storage_blob, query_blob); + ASSERT_NE(storage_blob, nullptr); + ASSERT_NE(query_blob, nullptr); + + // query blob should be aligned + unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; + ASSERT_EQ(address_alignment, 0); + + // They need to be processed by both processors. + ASSERT_EQ(((const float *)storage_blob)[0], initial_value + value_to_add_storage); + ASSERT_EQ(((const float *)query_blob)[0], initial_value + value_to_add_query); + + // the original blob should not change + ASSERT_NE(storage_blob, original_blob); + ASSERT_NE(query_blob, original_blob); + } + + // adding cosine preprocessor + auto cosine_preprocessor = new (allocator) CosinePreprocessor(allocator, dim); + multiPPContainer.addPreprocessor(cosine_preprocessor); + { + ProcessedBlobs processed_blobs = + multiPPContainer.preprocess(original_blob, sizeof(original_blob)); + const void *storage_blob = processed_blobs.getStorageBlob(); + const void *query_blob = processed_blobs.getQueryBlob(); + // blobs should point to a different memory slot + ASSERT_NE(storage_blob, query_blob); + // query memory should be aligned + unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; + ASSERT_EQ(address_alignment, 0); + // They need to be allocated and processed + ASSERT_NE(storage_blob, nullptr); + ASSERT_NE(query_blob, nullptr); + float expected_processed_storage[dim] = {initial_value + value_to_add_storage, + initial_value, initial_value, initial_value}; + float expected_processed_query[dim] = {initial_value + value_to_add_query, initial_value, + initial_value, initial_value}; + VecSim_Normalize(expected_processed_storage, dim, VecSimType_FLOAT32); + VecSim_Normalize(expected_processed_query, dim, VecSimType_FLOAT32); + ASSERT_EQ(((const float *)storage_blob)[0], expected_processed_storage[0]); + ASSERT_EQ(((const float *)query_blob)[0], expected_processed_query[0]); + // the original blob should not change + ASSERT_NE(storage_blob, original_blob); + ASSERT_NE(query_blob, original_blob); + } + // The preprocessors should be released by the preprocessors container. +} diff --git a/tests/unit/test_hnsw_tiered.cpp b/tests/unit/test_hnsw_tiered.cpp index 676424fd7..deb81c6aa 100644 --- a/tests/unit/test_hnsw_tiered.cpp +++ b/tests/unit/test_hnsw_tiered.cpp @@ -163,7 +163,8 @@ TYPED_TEST(HNSWTieredIndexTest, testIndexesAttributes) { dynamic_cast *>(bf_preprocessors) ->getPreprocessors(); const std::type_info &bf_pp_expected_type = typeid(CosinePreprocessor); - const std::type_info &bf_pp_actual_type = typeid(*pp_arr[0]); + PreprocessorInterface *bf_pp = pp_arr[0]; + const std::type_info &bf_pp_actual_type = typeid(*bf_pp); ASSERT_EQ(bf_pp_actual_type, bf_pp_expected_type); // hnsw - simple diff --git a/tests/unit/test_int8.cpp b/tests/unit/test_int8.cpp new file mode 100644 index 000000000..e298232ba --- /dev/null +++ b/tests/unit/test_int8.cpp @@ -0,0 +1,995 @@ +#include "gtest/gtest.h" +#include "VecSim/vec_sim.h" +#include "VecSim/algorithms/hnsw/hnsw_single.h" +#include "tests_utils.h" +#include "unit_test_utils.h" +#include "mock_thread_pool.h" +#include "VecSim/vec_sim_debug.h" +#include "VecSim/spaces/L2/L2.h" +#include "VecSim/spaces/IP/IP.h" + +class INT8Test : public ::testing::Test { +protected: + virtual void SetUp(HNSWParams ¶ms) { + FAIL() << "INT8Test::SetUp(HNSWParams) this method should be overriden"; + } + + virtual void SetUp(BFParams ¶ms) { + FAIL() << "INT8Test::SetUp(BFParams) this method should be overriden"; + } + + virtual void SetUp(TieredIndexParams &tiered_params) { + FAIL() << "INT8Test::SetUp(TieredIndexParams) this method should be overriden"; + } + + virtual void TearDown() { VecSimIndex_Free(index); } + + virtual const void *GetDataByInternalId(idType id) = 0; + + template + algo_t *CastIndex() { + return dynamic_cast(index); + } + + template + algo_t *CastIndex(VecSimIndex *vecsim_index) { + return dynamic_cast(vecsim_index); + } + + virtual HNSWIndex *CastToHNSW() { return CastIndex>(); } + + void PopulateRandomVector(int8_t *out_vec) { test_utils::populate_int8_vec(out_vec, dim); } + int PopulateRandomAndAddVector(size_t id, int8_t *out_vec) { + PopulateRandomVector(out_vec); + return VecSimIndex_AddVector(index, out_vec, id); + } + + virtual int GenerateAndAddVector(size_t id, int8_t value = 1) { + // use unit_test_utils.h + return ::GenerateAndAddVector(index, dim, id, value); + } + + void GenerateVector(int8_t *out_vec, int8_t value) { + // use unit_test_utils.h + return ::GenerateVector(out_vec, this->dim, value); + } + + virtual int GenerateRandomAndAddVector(size_t id) { + int8_t v[dim]; + PopulateRandomVector(v); + return VecSimIndex_AddVector(index, v, id); + } + + size_t GetValidVectorsCount() { + VecSimIndexInfo info = VecSimIndex_Info(index); + return info.commonInfo.indexLabelCount; + } + + template + void create_index_test(params_t index_params); + template + void element_size_test(params_t index_params); + template + void search_by_id_test(params_t index_params); + template + void search_by_score_test(params_t index_params); + template + void metrics_test(params_t index_params); + template + void search_empty_index_test(params_t index_params); + template + void test_override(params_t index_params); + template + void test_range_query(params_t index_params); + template + void test_batch_iterator_basic(params_t index_params); + template + VecSimIndexInfo test_info(params_t index_params); + template + void test_info_iterator(VecSimMetric metric); + template + void get_element_neighbors(params_t index_params); + + VecSimIndex *index; + size_t dim; +}; + +class INT8HNSWTest : public INT8Test { +protected: + virtual void SetUp(HNSWParams ¶ms) override { + params.type = VecSimType_INT8; + VecSimParams vecsim_params = CreateParams(params); + index = VecSimIndex_New(&vecsim_params); + dim = params.dim; + } + + virtual const void *GetDataByInternalId(idType id) override { + return CastIndex>()->getDataByInternalId(id); + } + + virtual HNSWIndex *CastToHNSW() override { + return CastIndex>(index); + } + + HNSWIndex *CastToHNSW(VecSimIndex *new_index) { + return CastIndex>(new_index); + } + + void test_info(bool is_multi); + void test_serialization(bool is_multi); +}; + +class INT8BruteForceTest : public INT8Test { +protected: + virtual void SetUp(BFParams ¶ms) override { + params.type = VecSimType_INT8; + VecSimParams vecsim_params = CreateParams(params); + index = VecSimIndex_New(&vecsim_params); + dim = params.dim; + } + + virtual const void *GetDataByInternalId(idType id) override { + return CastIndex>()->getDataByInternalId(id); + } + + virtual HNSWIndex *CastToHNSW() override { + ADD_FAILURE() << "INT8BruteForceTest::CastToHNSW() this method should not be called"; + return nullptr; + } + + void test_info(bool is_multi); +}; + +class INT8TieredTest : public INT8Test { +protected: + TieredIndexParams generate_tiered_params(HNSWParams &hnsw_params, size_t swap_job_threshold = 1, + size_t flat_buffer_limit = SIZE_MAX) { + hnsw_params.type = VecSimType_INT8; + vecsim_hnsw_params = CreateParams(hnsw_params); + TieredIndexParams tiered_params = { + .jobQueue = &mock_thread_pool.jobQ, + .jobQueueCtx = mock_thread_pool.ctx, + .submitCb = tieredIndexMock::submit_callback, + .flatBufferLimit = flat_buffer_limit, + .primaryIndexParams = &vecsim_hnsw_params, + .specificParams = {TieredHNSWParams{.swapJobThreshold = swap_job_threshold}}}; + return tiered_params; + } + + virtual void SetUp(TieredIndexParams &tiered_params) override { + VecSimParams params = CreateParams(tiered_params); + index = VecSimIndex_New(¶ms); + dim = tiered_params.primaryIndexParams->algoParams.hnswParams.dim; + + // Set the created tiered index in the index external context. + mock_thread_pool.ctx->index_strong_ref.reset(index); + } + + virtual void SetUp(HNSWParams &hnsw_params) override { + TieredIndexParams tiered_params = generate_tiered_params(hnsw_params); + SetUp(tiered_params); + } + + virtual void TearDown() override {} + + virtual const void *GetDataByInternalId(idType id) override { + return CastIndex>(CastToBruteForce()) + ->getDataByInternalId(id); + } + + virtual HNSWIndex *CastToHNSW() override { + auto tiered_index = dynamic_cast *>(index); + return tiered_index->getHNSWIndex(); + } + + virtual HNSWIndex_Single *CastToHNSWSingle() { + return CastIndex>(CastToHNSW()); + } + + VecSimIndexAbstract *CastToBruteForce() { + auto tiered_index = dynamic_cast *>(index); + return tiered_index->getFlatBufferIndex(); + } + + int GenerateRandomAndAddVector(size_t id) override { + int8_t v[dim]; + PopulateRandomVector(v); + int ret = VecSimIndex_AddVector(index, v, id); + mock_thread_pool.thread_iteration(); + return ret; + } + + int GenerateAndAddVector(size_t id, int8_t value) override { + // use unit_test_utils.h + int ret = INT8Test::GenerateAndAddVector(id, value); + mock_thread_pool.thread_iteration(); + return ret; + } + + void test_info(bool is_multi); + void test_info_iterator(VecSimMetric metric); + + VecSimParams vecsim_hnsw_params; + tieredIndexMock mock_thread_pool; +}; + +/* ---------------------------- Create index tests ---------------------------- */ + +template +void INT8Test::create_index_test(params_t index_params) { + SetUp(index_params); + + ASSERT_EQ(VecSimIndex_IndexSize(index), 0); + + int8_t vector[dim]; + this->PopulateRandomVector(vector); + VecSimIndex_AddVector(index, vector, 0); + + ASSERT_EQ(VecSimIndex_IndexSize(index), 1); + ASSERT_EQ(index->getDistanceFrom_Unsafe(0, vector), 0); + + ASSERT_NO_FATAL_FAILURE( + CompareVectors(static_cast(this->GetDataByInternalId(0)), vector, dim)); +} + +TEST_F(INT8HNSWTest, createIndex) { + HNSWParams params = {.dim = 40, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE(create_index_test(params)); + ASSERT_EQ(index->basicInfo().type, VecSimType_INT8); + ASSERT_EQ(index->basicInfo().algo, VecSimAlgo_HNSWLIB); +} + +TEST_F(INT8BruteForceTest, createIndex) { + BFParams params = {.dim = 40}; + EXPECT_NO_FATAL_FAILURE(create_index_test(params)); + ASSERT_EQ(index->basicInfo().type, VecSimType_INT8); + ASSERT_EQ(index->basicInfo().algo, VecSimAlgo_BF); +} + +TEST_F(INT8TieredTest, createIndex) { + HNSWParams params = {.dim = 40, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE(create_index_test(params)); + ASSERT_EQ(index->basicInfo().type, VecSimType_INT8); + ASSERT_EQ(index->basicInfo().isTiered, true); +} + +/* ---------------------------- Size Estimation tests ---------------------------- */ + +template +void INT8Test::element_size_test(params_t index_params) { + SetUp(index_params); + + // Estimate the memory delta of adding a single vector that requires a full new block. + size_t estimation = EstimateElementSize(index_params) * DEFAULT_BLOCK_SIZE; + size_t before = index->getAllocationSize(); + ASSERT_EQ(this->GenerateRandomAndAddVector(0), 1); + size_t actual = index->getAllocationSize() - before; + + // We check that the actual size is within 1% of the estimation. + ASSERT_GE(estimation, actual * 0.99); + ASSERT_LE(estimation, actual * 1.01); +} + +TEST_F(INT8HNSWTest, elementSizeEstimation) { + size_t M = 64; + + HNSWParams params = {.dim = 4, .M = M}; + EXPECT_NO_FATAL_FAILURE(element_size_test(params)); +} + +TEST_F(INT8BruteForceTest, elementSizeEstimation) { + BFParams params = {.dim = 4}; + EXPECT_NO_FATAL_FAILURE(element_size_test(params)); +} + +TEST_F(INT8TieredTest, elementSizeEstimation) { + size_t M = 64; + HNSWParams hnsw_params = {.dim = 4, .M = M}; + VecSimParams vecsim_hnsw_params = CreateParams(hnsw_params); + TieredIndexParams tiered_params = + test_utils::CreateTieredParams(vecsim_hnsw_params, this->mock_thread_pool); + EXPECT_NO_FATAL_FAILURE(element_size_test(tiered_params)); +} + +/* ---------------------------- Functionality tests ---------------------------- */ + +template +void INT8Test::search_by_id_test(params_t index_params) { + SetUp(index_params); + + size_t k = 11; + int8_t n = 100; + + for (int8_t i = 0; i < n; i++) { + this->GenerateAndAddVector(i, i); // {i, i, i, i} + } + ASSERT_EQ(VecSimIndex_IndexSize(index), n); + + int8_t query[dim]; + this->GenerateVector(query, 50); // {50, 50, 50, 50} + + // Vectors values are equal to the id, so the 11 closest vectors are 45, 46...50 + // (closest), 51...55 + static size_t expected_res_order[] = {45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55}; + auto verify_res = [&](size_t id, double score, size_t index) { + ASSERT_EQ(id, expected_res_order[index]); // results are sorted by ID + ASSERT_EQ(score, 4 * (50 - id) * (50 - id)); // L2 distance + }; + + runTopKSearchTest(index, query, k, verify_res, nullptr, BY_ID); +} + +TEST_F(INT8HNSWTest, searchByID) { + HNSWParams params = {.dim = 4, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE(search_by_id_test(params)); +} + +TEST_F(INT8BruteForceTest, searchByID) { + BFParams params = {.dim = 4}; + EXPECT_NO_FATAL_FAILURE(search_by_id_test(params)); +} + +TEST_F(INT8TieredTest, searchByID) { + HNSWParams params = {.dim = 4, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE(search_by_id_test(params)); +} + +template +void INT8Test::search_by_score_test(params_t index_params) { + SetUp(index_params); + + size_t k = 11; + size_t n = 100; + + for (size_t i = 0; i < n; i++) { + this->GenerateAndAddVector(i, i); // {i, i, i, i} + } + ASSERT_EQ(VecSimIndex_IndexSize(index), n); + + int8_t query[dim]; + this->GenerateVector(query, 50); // {50, 50, 50, 50} + + // Vectors values are equal to the id, so the 11 closest vectors are + // 45, 46...50 (closest), 51...55 + static size_t expected_res_order[] = {50, 49, 51, 48, 52, 47, 53, 46, 54, 45, 55}; + auto verify_res = [&](size_t id, double score, size_t index) { + ASSERT_EQ(id, expected_res_order[index]); + ASSERT_EQ(score, 4 * (50 - id) * (50 - id)); // L2 distance + }; + + // Search by score + runTopKSearchTest(index, query, k, verify_res); +} + +TEST_F(INT8HNSWTest, searchByScore) { + HNSWParams params = {.dim = 4, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE(search_by_score_test(params)); +} + +TEST_F(INT8BruteForceTest, searchByScore) { + BFParams params = {.dim = 4}; + EXPECT_NO_FATAL_FAILURE(search_by_score_test(params)); +} + +TEST_F(INT8TieredTest, searchByScore) { + HNSWParams params = {.dim = 4, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE(search_by_score_test(params)); +} + +template +void INT8Test::metrics_test(params_t index_params) { + SetUp(index_params); + size_t n = 10; + VecSimMetric metric = index_params.metric; + double expected_score = 0; + + auto verify_res = [&](size_t id, double score, size_t index) { + ASSERT_EQ(score, expected_score) << "failed at vector id:" << id; + }; + + for (size_t i = 0; i < n; i++) { + int8_t vector[dim]; + this->PopulateRandomAndAddVector(i, vector); + + if (metric == VecSimMetric_Cosine) { + // compare with the norm stored in the index vector + const int8_t *index_vector = static_cast(this->GetDataByInternalId(i)); + float index_vector_norm = *(reinterpret_cast(index_vector + dim)); + float vector_norm = spaces::IntegralType_ComputeNorm(vector, dim); + ASSERT_EQ(index_vector_norm, vector_norm) << "wrong vector norm for vector id:" << i; + } else if (metric == VecSimMetric_IP) { + expected_score = INT8_InnerProduct(vector, vector, dim); + } + + // query index with k = 1 expect to get the vector + runTopKSearchTest(index, vector, 1, verify_res); + ASSERT_EQ(VecSimIndex_IndexSize(index), i + 1); + } +} + +TEST_F(INT8HNSWTest, CosineTest) { + HNSWParams params = {.dim = 40, .metric = VecSimMetric_Cosine, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE(metrics_test(params)); +} +TEST_F(INT8HNSWTest, IPTest) { + HNSWParams params = {.dim = 40, .metric = VecSimMetric_IP, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE((metrics_test)(params)); +} +TEST_F(INT8HNSWTest, L2Test) { + HNSWParams params = {.dim = 40, .metric = VecSimMetric_L2, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE(metrics_test(params)); +} + +TEST_F(INT8BruteForceTest, CosineTest) { + BFParams params = {.dim = 40, .metric = VecSimMetric_Cosine}; + EXPECT_NO_FATAL_FAILURE(metrics_test(params)); +} +TEST_F(INT8BruteForceTest, IPTest) { + BFParams params = {.dim = 40, .metric = VecSimMetric_IP}; + EXPECT_NO_FATAL_FAILURE((metrics_test)(params)); +} +TEST_F(INT8BruteForceTest, L2Test) { + BFParams params = {.dim = 40, .metric = VecSimMetric_L2}; + EXPECT_NO_FATAL_FAILURE(metrics_test(params)); +} + +TEST_F(INT8TieredTest, CosineTest) { + HNSWParams params = {.dim = 40, .metric = VecSimMetric_Cosine, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE(metrics_test(params)); +} +TEST_F(INT8TieredTest, IPTest) { + HNSWParams params = {.dim = 40, .metric = VecSimMetric_IP, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE((metrics_test)(params)); +} +TEST_F(INT8TieredTest, L2Test) { + HNSWParams params = {.dim = 40, .metric = VecSimMetric_L2, .M = 16, .efConstruction = 200}; + EXPECT_NO_FATAL_FAILURE(metrics_test(params)); +} + +template +void INT8Test::search_empty_index_test(params_t params) { + size_t n = 100; + size_t k = 11; + + SetUp(params); + ASSERT_EQ(VecSimIndex_IndexSize(index), 0); + + int8_t query[dim]; + this->GenerateVector(query, 50); // {50, 50, 50, 50} + + // We do not expect any results. + VecSimQueryReply *res = VecSimIndex_TopKQuery(index, query, k, NULL, BY_SCORE); + ASSERT_EQ(VecSimQueryReply_Len(res), 0); + VecSimQueryReply_Iterator *it = VecSimQueryReply_GetIterator(res); + ASSERT_EQ(VecSimQueryReply_IteratorNext(it), nullptr); + VecSimQueryReply_IteratorFree(it); + VecSimQueryReply_Free(res); + + res = VecSimIndex_RangeQuery(index, query, 1.0, NULL, BY_SCORE); + ASSERT_EQ(VecSimQueryReply_Len(res), 0); + VecSimQueryReply_Free(res); + + // Add some vectors and remove them all from index, so it will be empty again. + for (size_t i = 0; i < n; i++) { + this->GenerateAndAddVector(i); + } + ASSERT_EQ(VecSimIndex_IndexSize(index), n); + for (size_t i = 0; i < n; i++) { + VecSimIndex_DeleteVector(index, i); + } + // vectors marked as deleted will be included in VecSimIndex_IndexSize + ASSERT_EQ(GetValidVectorsCount(), 0); + + // Again - we do not expect any results. + res = VecSimIndex_TopKQuery(index, query, k, NULL, BY_SCORE); + ASSERT_EQ(VecSimQueryReply_Len(res), 0); + it = VecSimQueryReply_GetIterator(res); + ASSERT_EQ(VecSimQueryReply_IteratorNext(it), nullptr); + VecSimQueryReply_IteratorFree(it); + VecSimQueryReply_Free(res); + + res = VecSimIndex_RangeQuery(index, query, 1.0, NULL, BY_SCORE); + ASSERT_EQ(VecSimQueryReply_Len(res), 0); + VecSimQueryReply_Free(res); +} + +TEST_F(INT8HNSWTest, SearchEmptyIndex) { + HNSWParams params = {.dim = 4, .initialCapacity = 0}; + EXPECT_NO_FATAL_FAILURE(search_empty_index_test(params)); +} + +TEST_F(INT8BruteForceTest, SearchEmptyIndex) { + BFParams params = {.dim = 4, .initialCapacity = 0}; + EXPECT_NO_FATAL_FAILURE(search_empty_index_test(params)); +} + +TEST_F(INT8TieredTest, SearchEmptyIndex) { + HNSWParams params = {.dim = 4, .initialCapacity = 0}; + EXPECT_NO_FATAL_FAILURE(search_empty_index_test(params)); +} + +template +void INT8Test::test_override(params_t params) { + size_t n = 50; + size_t new_n = 120; + SetUp(params); + + // Insert n vectors. + for (size_t i = 0; i < n; i++) { + ASSERT_EQ(GenerateAndAddVector(i, i), 1); + } + ASSERT_EQ(VecSimIndex_IndexSize(index), n); + + // Override n vectors, the first 100 will be overwritten (deleted first). + for (size_t i = 0; i < n; i++) { + ASSERT_EQ(this->GenerateAndAddVector(i, i), 0); + } + + // Add up to new_n vectors. + for (size_t i = n; i < new_n; i++) { + ASSERT_EQ(this->GenerateAndAddVector(i, i), 1); + } + + int8_t query[dim]; + this->GenerateVector(query, new_n); + + // Vectors values equals their id, so we expect the larger the id the closest it will be to the + // query. + auto verify_res = [&](size_t id, double score, size_t index) { + ASSERT_EQ(id, new_n - 1 - index) << "id: " << id << " score: " << score; + float diff = new_n - id; + float exp_score = 4 * diff * diff; + ASSERT_EQ(score, exp_score) << "id: " << id << " score: " << score; + }; + runTopKSearchTest(index, query, 300, verify_res); +} + +TEST_F(INT8HNSWTest, Override) { + HNSWParams params = { + .dim = 4, .initialCapacity = 100, .M = 8, .efConstruction = 20, .efRuntime = 250}; + EXPECT_NO_FATAL_FAILURE(test_override(params)); +} + +TEST_F(INT8BruteForceTest, Override) { + BFParams params = {.dim = 4, .initialCapacity = 100}; + EXPECT_NO_FATAL_FAILURE(test_override(params)); +} + +TEST_F(INT8TieredTest, Override) { + HNSWParams params = { + .dim = 4, .initialCapacity = 100, .M = 8, .efConstruction = 20, .efRuntime = 250}; + EXPECT_NO_FATAL_FAILURE(test_override(params)); +} + +template +void INT8Test::test_range_query(params_t params) { + size_t n = 100; + SetUp(params); + + int8_t pivot_value = 1; + int8_t pivot_vec[dim]; + this->GenerateVector(pivot_vec, pivot_value); + + int8_t radius = 20; + std::mt19937 gen(42); + std::uniform_int_distribution dis(pivot_value - radius, pivot_value + radius); + + // insert 20 vectors near a pivot vector. + size_t n_close = 20; + for (size_t i = 0; i < n_close; i++) { + int8_t random_number = static_cast(dis(gen)); + this->GenerateAndAddVector(i, random_number); + } + + int8_t max_vec[dim]; + GenerateVector(max_vec, pivot_value + radius); + float max_dist = INT8_L2Sqr(pivot_vec, max_vec, dim); + + // Add more vectors far from the pivot vector + for (size_t i = n_close; i < n; i++) { + int8_t random_number = static_cast(dis(gen)); + GenerateAndAddVector(i, 50 + random_number); + } + ASSERT_EQ(VecSimIndex_IndexSize(index), n); + + auto verify_res_by_score = [&](size_t id, double score, size_t index) { + ASSERT_LE(id, n_close - 1) << "score: " << score; + ASSERT_LE(score, max_dist); + }; + size_t expected_num_results = n_close; + + runRangeQueryTest(index, pivot_vec, max_dist, verify_res_by_score, expected_num_results, + BY_SCORE); +} + +TEST_F(INT8HNSWTest, rangeQuery) { + HNSWParams params = {.dim = 4}; + EXPECT_NO_FATAL_FAILURE(test_range_query(params)); +} + +TEST_F(INT8BruteForceTest, rangeQuery) { + BFParams params = {.dim = 4}; + EXPECT_NO_FATAL_FAILURE(test_range_query(params)); +} + +TEST_F(INT8TieredTest, rangeQuery) { + HNSWParams params = {.dim = 4}; + EXPECT_NO_FATAL_FAILURE(test_range_query(params)); +} + +/* ---------------------------- Batch iterator tests ---------------------------- */ + +template +void INT8Test::test_batch_iterator_basic(params_t params) { + SetUp(params); + size_t n = 100; + + // For every i, add the vector (i,i,i,i) under the label i. + for (size_t i = 0; i < n; i++) { + ASSERT_EQ(this->GenerateAndAddVector(i, i), 1); + } + + ASSERT_EQ(VecSimIndex_IndexSize(index), n); + + // Query for (n,n,n,n) vector (recall that n-1 is the largest id in te index). + int8_t query[dim]; + GenerateVector(query, n); + + VecSimBatchIterator *batchIterator = VecSimBatchIterator_New(index, query, nullptr); + size_t iteration_num = 0; + + // Get the 5 vectors whose ids are the maximal among those that hasn't been returned yet + // in every iteration. The results order should be sorted by their score (distance from the + // query vector), which means sorted from the largest id to the lowest. + size_t n_res = 5; + while (VecSimBatchIterator_HasNext(batchIterator)) { + std::vector expected_ids(n_res); + for (size_t i = 0; i < n_res; i++) { + expected_ids[i] = (n - iteration_num * n_res - i - 1); + } + auto verify_res = [&](size_t id, double score, size_t index) { + ASSERT_EQ(expected_ids[index], id) + << "iteration_num: " << iteration_num << " index: " << index << " score: " << score; + }; + runBatchIteratorSearchTest(batchIterator, n_res, verify_res); + iteration_num++; + } + ASSERT_EQ(iteration_num, n / n_res); + VecSimBatchIterator_Free(batchIterator); +} + +TEST_F(INT8HNSWTest, BatchIteratorBasic) { + HNSWParams params = {.dim = 4, .M = 8, .efConstruction = 20, .efRuntime = 100}; + EXPECT_NO_FATAL_FAILURE(test_batch_iterator_basic(params)); +} + +TEST_F(INT8BruteForceTest, BatchIteratorBasic) { + BFParams params = {.dim = 4}; + EXPECT_NO_FATAL_FAILURE(test_batch_iterator_basic(params)); +} + +TEST_F(INT8TieredTest, BatchIteratorBasic) { + HNSWParams params = {.dim = 4, .M = 8, .efConstruction = 20, .efRuntime = 100}; + EXPECT_NO_FATAL_FAILURE(test_batch_iterator_basic(params)); +} + +/* ---------------------------- Info tests ---------------------------- */ + +template +VecSimIndexInfo INT8Test::test_info(params_t params) { + SetUp(params); + VecSimIndexInfo info = VecSimIndex_Info(index); + EXPECT_EQ(info.commonInfo.basicInfo.dim, params.dim); + EXPECT_EQ(info.commonInfo.basicInfo.isMulti, params.multi); + EXPECT_EQ(info.commonInfo.basicInfo.type, VecSimType_INT8); + EXPECT_EQ(info.commonInfo.basicInfo.blockSize, DEFAULT_BLOCK_SIZE); + EXPECT_EQ(info.commonInfo.indexSize, 0); + EXPECT_EQ(info.commonInfo.indexLabelCount, 0); + EXPECT_EQ(info.commonInfo.memory, index->getAllocationSize()); + EXPECT_EQ(info.commonInfo.basicInfo.metric, VecSimMetric_L2); + + // Validate that basic info returns the right restricted info as well. + VecSimIndexBasicInfo s_info = VecSimIndex_BasicInfo(index); + EXPECT_EQ(info.commonInfo.basicInfo.algo, s_info.algo); + EXPECT_EQ(info.commonInfo.basicInfo.dim, s_info.dim); + EXPECT_EQ(info.commonInfo.basicInfo.blockSize, s_info.blockSize); + EXPECT_EQ(info.commonInfo.basicInfo.type, s_info.type); + EXPECT_EQ(info.commonInfo.basicInfo.isMulti, s_info.isMulti); + EXPECT_EQ(info.commonInfo.basicInfo.type, s_info.type); + EXPECT_EQ(info.commonInfo.basicInfo.isTiered, s_info.isTiered); + + return info; +} + +void INT8HNSWTest::test_info(bool is_multi) { + HNSWParams params = {.dim = 128, .multi = is_multi}; + VecSimIndexInfo info = INT8Test::test_info(params); + ASSERT_EQ(info.commonInfo.basicInfo.algo, VecSimAlgo_HNSWLIB); + + ASSERT_EQ(info.hnswInfo.M, HNSW_DEFAULT_M); + ASSERT_EQ(info.hnswInfo.efConstruction, HNSW_DEFAULT_EF_C); + ASSERT_EQ(info.hnswInfo.efRuntime, HNSW_DEFAULT_EF_RT); + ASSERT_DOUBLE_EQ(info.hnswInfo.epsilon, HNSW_DEFAULT_EPSILON); +} +TEST_F(INT8HNSWTest, testInfoSingle) { test_info(false); } +TEST_F(INT8HNSWTest, testInfoMulti) { test_info(true); } + +void INT8BruteForceTest::test_info(bool is_multi) { + BFParams params = {.dim = 128, .multi = is_multi}; + VecSimIndexInfo info = INT8Test::test_info(params); + ASSERT_EQ(info.commonInfo.basicInfo.algo, VecSimAlgo_BF); +} + +TEST_F(INT8BruteForceTest, testInfoSingle) { test_info(false); } +TEST_F(INT8BruteForceTest, testInfoMulti) { test_info(true); } + +void INT8TieredTest::test_info(bool is_multi) { + size_t bufferLimit = SIZE_MAX; + HNSWParams hnsw_params = {.dim = 128, .multi = is_multi}; + + VecSimIndexInfo info = INT8Test::test_info(hnsw_params); + ASSERT_EQ(info.commonInfo.basicInfo.algo, VecSimAlgo_HNSWLIB); + VecSimIndexInfo frontendIndexInfo = CastToBruteForce()->info(); + VecSimIndexInfo backendIndexInfo = CastToHNSW()->info(); + + compareCommonInfo(info.tieredInfo.frontendCommonInfo, frontendIndexInfo.commonInfo); + compareFlatInfo(info.tieredInfo.bfInfo, frontendIndexInfo.bfInfo); + compareCommonInfo(info.tieredInfo.backendCommonInfo, backendIndexInfo.commonInfo); + compareHNSWInfo(info.tieredInfo.backendInfo.hnswInfo, backendIndexInfo.hnswInfo); + + EXPECT_EQ(info.commonInfo.memory, info.tieredInfo.management_layer_memory + + backendIndexInfo.commonInfo.memory + + frontendIndexInfo.commonInfo.memory); + EXPECT_EQ(info.tieredInfo.backgroundIndexing, false); + EXPECT_EQ(info.tieredInfo.bufferLimit, bufferLimit); + EXPECT_EQ(info.tieredInfo.specificTieredBackendInfo.hnswTieredInfo.pendingSwapJobsThreshold, 1); + + INT8Test::GenerateAndAddVector(1, 1); + info = index->info(); + + EXPECT_EQ(info.commonInfo.indexSize, 1); + EXPECT_EQ(info.commonInfo.indexLabelCount, 1); + EXPECT_EQ(info.tieredInfo.backendCommonInfo.indexSize, 0); + EXPECT_EQ(info.tieredInfo.backendCommonInfo.indexLabelCount, 0); + EXPECT_EQ(info.tieredInfo.frontendCommonInfo.indexSize, 1); + EXPECT_EQ(info.tieredInfo.frontendCommonInfo.indexLabelCount, 1); + EXPECT_EQ(info.commonInfo.memory, info.tieredInfo.management_layer_memory + + info.tieredInfo.backendCommonInfo.memory + + info.tieredInfo.frontendCommonInfo.memory); + EXPECT_EQ(info.tieredInfo.backgroundIndexing, true); + + mock_thread_pool.thread_iteration(); + info = index->info(); + + EXPECT_EQ(info.commonInfo.indexSize, 1); + EXPECT_EQ(info.commonInfo.indexLabelCount, 1); + EXPECT_EQ(info.tieredInfo.backendCommonInfo.indexSize, 1); + EXPECT_EQ(info.tieredInfo.backendCommonInfo.indexLabelCount, 1); + EXPECT_EQ(info.tieredInfo.frontendCommonInfo.indexSize, 0); + EXPECT_EQ(info.tieredInfo.frontendCommonInfo.indexLabelCount, 0); + EXPECT_EQ(info.commonInfo.memory, info.tieredInfo.management_layer_memory + + info.tieredInfo.backendCommonInfo.memory + + info.tieredInfo.frontendCommonInfo.memory); + EXPECT_EQ(info.tieredInfo.backgroundIndexing, false); + + if (is_multi) { + INT8Test::GenerateAndAddVector(1, 1); + info = index->info(); + + EXPECT_EQ(info.commonInfo.indexSize, 2); + EXPECT_EQ(info.commonInfo.indexLabelCount, 1); + EXPECT_EQ(info.tieredInfo.backendCommonInfo.indexSize, 1); + EXPECT_EQ(info.tieredInfo.backendCommonInfo.indexLabelCount, 1); + EXPECT_EQ(info.tieredInfo.frontendCommonInfo.indexSize, 1); + EXPECT_EQ(info.tieredInfo.frontendCommonInfo.indexLabelCount, 1); + EXPECT_EQ(info.commonInfo.memory, info.tieredInfo.management_layer_memory + + info.tieredInfo.backendCommonInfo.memory + + info.tieredInfo.frontendCommonInfo.memory); + EXPECT_EQ(info.tieredInfo.backgroundIndexing, true); + } + + VecSimIndex_DeleteVector(index, 1); + info = index->info(); + + EXPECT_EQ(info.commonInfo.indexSize, 0); + EXPECT_EQ(info.commonInfo.indexLabelCount, 0); + EXPECT_EQ(info.tieredInfo.backendCommonInfo.indexSize, 0); + EXPECT_EQ(info.tieredInfo.backendCommonInfo.indexLabelCount, 0); + EXPECT_EQ(info.tieredInfo.frontendCommonInfo.indexSize, 0); + EXPECT_EQ(info.tieredInfo.frontendCommonInfo.indexLabelCount, 0); + EXPECT_EQ(info.commonInfo.memory, info.tieredInfo.management_layer_memory + + info.tieredInfo.backendCommonInfo.memory + + info.tieredInfo.frontendCommonInfo.memory); + EXPECT_EQ(info.tieredInfo.backgroundIndexing, false); +} + +TEST_F(INT8TieredTest, testInfoSingle) { test_info(false); } +TEST_F(INT8TieredTest, testInfoMulti) { test_info(true); } + +template +void INT8Test::test_info_iterator(VecSimMetric metric) { + params_t params = {.dim = 128, .metric = metric}; + SetUp(params); + VecSimIndexInfo info = VecSimIndex_Info(index); + VecSimInfoIterator *infoIter = VecSimIndex_InfoIterator(index); + VecSimAlgo algo = info.commonInfo.basicInfo.algo; + if (algo == VecSimAlgo_HNSWLIB) { + compareHNSWIndexInfoToIterator(info, infoIter); + } else if (algo == VecSimAlgo_BF) { + compareFlatIndexInfoToIterator(info, infoIter); + } + VecSimInfoIterator_Free(infoIter); +} + +TEST_F(INT8BruteForceTest, InfoIteratorCosine) { + test_info_iterator(VecSimMetric_Cosine); +} +TEST_F(INT8BruteForceTest, InfoIteratorIP) { test_info_iterator(VecSimMetric_IP); } +TEST_F(INT8BruteForceTest, InfoIteratorL2) { test_info_iterator(VecSimMetric_L2); } +TEST_F(INT8HNSWTest, InfoIteratorCosine) { test_info_iterator(VecSimMetric_Cosine); } +TEST_F(INT8HNSWTest, InfoIteratorIP) { test_info_iterator(VecSimMetric_IP); } +TEST_F(INT8HNSWTest, InfoIteratorL2) { test_info_iterator(VecSimMetric_L2); } + +void INT8TieredTest::test_info_iterator(VecSimMetric metric) { + size_t n = 100; + size_t d = 128; + HNSWParams params = {.dim = d, .metric = metric, .initialCapacity = n}; + SetUp(params); + VecSimIndexInfo info = VecSimIndex_Info(index); + VecSimInfoIterator *infoIter = VecSimIndex_InfoIterator(index); + VecSimIndexInfo frontendIndexInfo = CastToBruteForce()->info(); + VecSimIndexInfo backendIndexInfo = CastToHNSW()->info(); + VecSimInfoIterator_Free(infoIter); +} + +TEST_F(INT8TieredTest, InfoIteratorCosine) { test_info_iterator(VecSimMetric_Cosine); } +TEST_F(INT8TieredTest, InfoIteratorIP) { test_info_iterator(VecSimMetric_IP); } +TEST_F(INT8TieredTest, InfoIteratorL2) { test_info_iterator(VecSimMetric_L2); } + +/* ---------------------------- HNSW specific tests ---------------------------- */ + +void INT8HNSWTest::test_serialization(bool is_multi) { + size_t dim = 4; + size_t n = 1001; + size_t n_labels[] = {n, 100}; + size_t M = 8; + size_t ef = 10; + double epsilon = 0.004; + size_t blockSize = 20; + std::string multiToString[] = {"single", "multi_100labels_"}; + + HNSWParams params{.type = VecSimType_INT8, + .dim = dim, + .metric = VecSimMetric_Cosine, + .multi = is_multi, + .initialCapacity = n, + .blockSize = blockSize, + .M = M, + .efConstruction = ef, + .efRuntime = ef, + .epsilon = epsilon}; + SetUp(params); + + auto *hnsw_index = this->CastToHNSW(); + + int8_t data[n * dim]; + + for (size_t i = 0; i < n * dim; i += dim) { + test_utils::populate_int8_vec(data + i, dim, i); + } + + for (size_t j = 0; j < n; ++j) { + VecSimIndex_AddVector(index, data + dim * j, j % n_labels[is_multi]); + } + + auto file_name = std::string(getenv("ROOT")) + "/tests/unit/1k-d4-L2-M8-ef_c10_" + + VecSimType_ToString(VecSimType_INT8) + "_" + multiToString[is_multi] + + ".hnsw_current_version"; + + // Save the index with the default version (V3). + hnsw_index->saveIndex(file_name); + + // Fetch info after saving, as memory size change during saving. + VecSimIndexInfo info = VecSimIndex_Info(index); + ASSERT_EQ(info.commonInfo.basicInfo.algo, VecSimAlgo_HNSWLIB); + ASSERT_EQ(info.hnswInfo.M, M); + ASSERT_EQ(info.hnswInfo.efConstruction, ef); + ASSERT_EQ(info.hnswInfo.efRuntime, ef); + ASSERT_EQ(info.commonInfo.indexSize, n); + ASSERT_EQ(info.commonInfo.basicInfo.metric, VecSimMetric_Cosine); + ASSERT_EQ(info.commonInfo.basicInfo.type, VecSimType_INT8); + ASSERT_EQ(info.commonInfo.basicInfo.dim, dim); + ASSERT_EQ(info.commonInfo.indexLabelCount, n_labels[is_multi]); + + // Load the index from the file. + VecSimIndex *serialized_index = HNSWFactory::NewIndex(file_name); + auto *serialized_hnsw_index = this->CastToHNSW(serialized_index); + + // Verify that the index was loaded as expected. + ASSERT_TRUE(serialized_hnsw_index->checkIntegrity().valid_state); + + VecSimIndexInfo info2 = VecSimIndex_Info(serialized_index); + ASSERT_EQ(info2.commonInfo.basicInfo.algo, VecSimAlgo_HNSWLIB); + ASSERT_EQ(info2.hnswInfo.M, M); + ASSERT_EQ(info2.commonInfo.basicInfo.isMulti, is_multi); + ASSERT_EQ(info2.commonInfo.basicInfo.blockSize, blockSize); + ASSERT_EQ(info2.hnswInfo.efConstruction, ef); + ASSERT_EQ(info2.hnswInfo.efRuntime, ef); + ASSERT_EQ(info2.commonInfo.indexSize, n); + ASSERT_EQ(info2.commonInfo.basicInfo.metric, VecSimMetric_Cosine); + ASSERT_EQ(info2.commonInfo.basicInfo.type, VecSimType_INT8); + ASSERT_EQ(info2.commonInfo.basicInfo.dim, dim); + ASSERT_EQ(info2.commonInfo.indexLabelCount, n_labels[is_multi]); + ASSERT_EQ(info2.hnswInfo.epsilon, epsilon); + + // Check the functionality of the loaded index. + + int8_t new_vec[dim]; + this->PopulateRandomVector(new_vec); + VecSimIndex_AddVector(serialized_index, new_vec, n); + auto verify_res = [&](size_t id, double score, size_t index) { + ASSERT_EQ(id, n) << "score: " << score; + ASSERT_NEAR(score, 0.0, 1e-7); + }; + runTopKSearchTest(serialized_index, new_vec, 1, verify_res); + VecSimIndex_DeleteVector(serialized_index, 1); + + size_t n_per_label = n / n_labels[is_multi]; + ASSERT_TRUE(serialized_hnsw_index->checkIntegrity().valid_state); + ASSERT_EQ(VecSimIndex_IndexSize(serialized_index), n + 1 - n_per_label); + + // Clean up. + remove(file_name.c_str()); + VecSimIndex_Free(serialized_index); +} + +TEST_F(INT8HNSWTest, SerializationCurrentVersion) { test_serialization(false); } + +TEST_F(INT8HNSWTest, SerializationCurrentVersionMulti) { test_serialization(true); } + +template +void INT8Test::get_element_neighbors(params_t params) { + size_t n = 0; + + SetUp(params); + auto *hnsw_index = CastToHNSW(); + + // Add vectors until we have at least 2 vectors at level 1. + size_t vectors_in_higher_levels = 0; + while (vectors_in_higher_levels < 2) { + GenerateAndAddVector(n, n); + if (hnsw_index->getGraphDataByInternalId(n)->toplevel > 0) { + vectors_in_higher_levels++; + } + n++; + } + ASSERT_GE(n, 1) << "n: " << n; + + // Go over all vectors and validate that the getElementNeighbors debug command returns the + // neighbors properly. + for (size_t id = 0; id < n; id++) { + ElementLevelData &cur = hnsw_index->getElementLevelData(id, 0); + int **neighbors_output; + VecSimDebug_GetElementNeighborsInHNSWGraph(index, id, &neighbors_output); + auto graph_data = hnsw_index->getGraphDataByInternalId(id); + for (size_t l = 0; l <= graph_data->toplevel; l++) { + auto &level_data = hnsw_index->getElementLevelData(graph_data, l); + auto &neighbours = neighbors_output[l]; + ASSERT_EQ(neighbours[0], level_data.numLinks); + for (size_t j = 1; j <= neighbours[0]; j++) { + ASSERT_EQ(neighbours[j], level_data.links[j - 1]); + } + } + VecSimDebug_ReleaseElementNeighborsInHNSWGraph(neighbors_output); + } +} + +TEST_F(INT8HNSWTest, getElementNeighbors) { + HNSWParams params = {.dim = 4, .M = 20}; + get_element_neighbors(params); +} + +TEST_F(INT8TieredTest, getElementNeighbors) { + HNSWParams params = {.dim = 4, .M = 20}; + get_element_neighbors(params); +} diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index 9931d318a..e554c88ef 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -242,15 +242,15 @@ TEST_F(SpacesTest, int8_ip_no_optimization_func_test) { TEST_F(SpacesTest, int8_Cosine_no_optimization_func_test) { size_t dim = 4; // create a vector with extra space for the norm - int8_t *v1 = new int8_t[dim + sizeof(float)]; - int8_t *v2 = new int8_t[dim + sizeof(float)]; + int8_t v1[dim + sizeof(float)]; + int8_t v2[dim + sizeof(float)]; test_utils::populate_int8_vec(v1, dim, 123); test_utils::populate_int8_vec(v2, dim, 123); // write the norm at the end of the vector - *(float *)(v1 + dim) = test_utils::compute_norm(v1, dim); - *(float *)(v2 + dim) = test_utils::compute_norm(v2, dim); + *(float *)(v1 + dim) = test_utils::integral_compute_norm(v1, dim); + *(float *)(v2 + dim) = test_utils::integral_compute_norm(v2, dim); float dist = INT8_Cosine((const void *)v1, (const void *)v2, dim); ASSERT_NEAR(dist, 0.0, 0.000001); @@ -917,8 +917,8 @@ class INT8SpacesOptimizationTest : public testing::TestWithParam {}; TEST_P(INT8SpacesOptimizationTest, INT8L2SqrTest) { auto optimization = cpu_features::GetX86Info().features; size_t dim = GetParam(); - int8_t *v1 = new int8_t[dim]; - int8_t *v2 = new int8_t[dim]; + int8_t v1[dim]; + int8_t v2[dim]; test_utils::populate_int8_vec(v1, dim, 123); test_utils::populate_int8_vec(v2, dim, 1234); @@ -953,8 +953,8 @@ TEST_P(INT8SpacesOptimizationTest, INT8L2SqrTest) { TEST_P(INT8SpacesOptimizationTest, INT8InnerProductTest) { auto optimization = cpu_features::GetX86Info().features; size_t dim = GetParam(); - int8_t *v1 = new int8_t[dim]; - int8_t *v2 = new int8_t[dim]; + int8_t v1[dim]; + int8_t v2[dim]; test_utils::populate_int8_vec(v1, dim, 123); test_utils::populate_int8_vec(v2, dim, 1234); @@ -990,14 +990,14 @@ TEST_P(INT8SpacesOptimizationTest, INT8InnerProductTest) { TEST_P(INT8SpacesOptimizationTest, INT8CosineTest) { auto optimization = cpu_features::GetX86Info().features; size_t dim = GetParam(); - int8_t *v1 = new int8_t[dim + sizeof(float)]; - int8_t *v2 = new int8_t[dim + sizeof(float)]; + int8_t v1[dim + sizeof(float)]; + int8_t v2[dim + sizeof(float)]; test_utils::populate_int8_vec(v1, dim, 123); test_utils::populate_int8_vec(v2, dim, 1234); // write the norm at the end of the vector - *(float *)(v1 + dim) = test_utils::compute_norm(v1, dim); - *(float *)(v2 + dim) = test_utils::compute_norm(v2, dim); + *(float *)(v1 + dim) = test_utils::integral_compute_norm(v1, dim); + *(float *)(v2 + dim) = test_utils::integral_compute_norm(v2, dim); dist_func_t arch_opt_func; float baseline = INT8_Cosine(v1, v2, dim); diff --git a/tests/unit/unit_test_utils.cpp b/tests/unit/unit_test_utils.cpp index 89973d19d..41d8dbcb5 100644 --- a/tests/unit/unit_test_utils.cpp +++ b/tests/unit/unit_test_utils.cpp @@ -46,6 +46,7 @@ VecSimQueryParams CreateQueryParams(const HNSWRuntimeParams &RuntimeParams) { static bool is_async_index(VecSimIndex *index) { return dynamic_cast *>(index) != nullptr || + dynamic_cast *>(index) != nullptr || dynamic_cast *>(index) != nullptr; } @@ -376,3 +377,69 @@ size_t getLabelsLookupNodeSize() { size_t memory_after = allocator->getAllocationSize(); return memory_after - memory_before; } +namespace test_utils { +size_t CalcVectorDataSize(VecSimIndex *index, VecSimType data_type) { + switch (data_type) { + case VecSimType_FLOAT32: { + VecSimIndexAbstract *abs_index = + dynamic_cast *>(index); + assert(abs_index && + "dynamic_cast failed: can't convert index to VecSimIndexAbstract"); + return abs_index->getDataSize(); + } + case VecSimType_FLOAT64: { + VecSimIndexAbstract *abs_index = + dynamic_cast *>(index); + assert(abs_index && + "dynamic_cast failed: can't convert index to VecSimIndexAbstract"); + return abs_index->getDataSize(); + } + case VecSimType_BFLOAT16: { + VecSimIndexAbstract *abs_index = + dynamic_cast *>(index); + assert(abs_index && "dynamic_cast failed: can't convert index to " + "VecSimIndexAbstract"); + return abs_index->getDataSize(); + } + case VecSimType_FLOAT16: { + VecSimIndexAbstract *abs_index = + dynamic_cast *>(index); + assert(abs_index && "dynamic_cast failed: can't convert index to " + "VecSimIndexAbstract"); + return abs_index->getDataSize(); + } + case VecSimType_INT8: { + VecSimIndexAbstract *abs_index = + dynamic_cast *>(index); + assert(abs_index && + "dynamic_cast failed: can't convert index to VecSimIndexAbstract"); + return abs_index->getDataSize(); + } + default: + return 0; + } +} + +TieredIndexParams CreateTieredParams(VecSimParams &primary_params, + tieredIndexMock &mock_thread_pool) { + TieredIndexParams tiered_params = {.jobQueue = &mock_thread_pool.jobQ, + .jobQueueCtx = mock_thread_pool.ctx, + .submitCb = tieredIndexMock::submit_callback, + .flatBufferLimit = SIZE_MAX, + .primaryIndexParams = &primary_params, + .specificParams = {TieredHNSWParams{.swapJobThreshold = 0}}}; + + return tiered_params; +} + +VecSimIndex *CreateNewTieredHNSWIndex(const HNSWParams &hnsw_params, + tieredIndexMock &mock_thread_pool) { + VecSimParams primary_params = CreateParams(hnsw_params); + auto tiered_params = CreateTieredParams(primary_params, mock_thread_pool); + VecSimParams params = CreateParams(tiered_params); + VecSimIndex *index = VecSimIndex_New(¶ms); + mock_thread_pool.ctx->index_strong_ref.reset(index); + + return index; +} +} // namespace test_utils diff --git a/tests/unit/unit_test_utils.h b/tests/unit/unit_test_utils.h index 54478cfd7..cafde7552 100644 --- a/tests/unit/unit_test_utils.h +++ b/tests/unit/unit_test_utils.h @@ -13,6 +13,7 @@ #include "VecSim/vec_sim.h" #include "VecSim/algorithms/hnsw/hnsw_tiered.h" +#include "mock_thread_pool.h" #include "gtest/gtest.h" // IndexType is used to define indices unit tests @@ -99,6 +100,11 @@ inline VecSimIndex *CreateNewIndex(IndexParams &index_params, VecSimType type, return VecSimIndex_New(¶ms); } +TieredIndexParams CreateTieredParams(VecSimParams &primary_params, + tieredIndexMock &mock_thread_pool); +VecSimIndex *CreateNewTieredHNSWIndex(const HNSWParams &hnsw_params, + tieredIndexMock &mock_thread_pool); + extern VecsimQueryType query_types[4]; } // namespace test_utils @@ -162,6 +168,16 @@ inline double GetInfVal(VecSimType type) { throw std::invalid_argument("This type is not supported"); } } +// TODO: Move all test_utils to this namespace +namespace test_utils { +size_t CalcVectorDataSize(VecSimIndex *index, VecSimType data_type); + +template +TieredHNSWIndex *cast_to_tiered_index(VecSimIndex *index) { + return dynamic_cast *>(index); +} + +} // namespace test_utils // Test a specific exception type is thrown and prints the right message. #define ASSERT_EXCEPTION_MESSAGE(VALUE, EXCEPTION_TYPE, MESSAGE) \ diff --git a/tests/utils/tests_utils.h b/tests/utils/tests_utils.h index 31dc3d9ef..0bf8bca53 100644 --- a/tests/utils/tests_utils.h +++ b/tests/utils/tests_utils.h @@ -2,6 +2,7 @@ #include #include +#include "VecSim/spaces/normalize/compute_norm.h" namespace test_utils { @@ -19,13 +20,9 @@ static void populate_int8_vec(int8_t *v, size_t dim, int seed = 1234) { } } -// TODO: replace with normalize function from VecSim -float compute_norm(const int8_t *vec, size_t dim) { - int norm = 0; - for (size_t i = 0; i < dim; i++) { - norm += vec[i] * vec[i]; - } - return sqrt(norm); +template +float integral_compute_norm(const datatype *vec, size_t dim) { + return spaces::IntegralType_ComputeNorm(vec, dim); } } // namespace test_utils