From 709b9984b2405464edaa2a7c71204a3ff112b3c4 Mon Sep 17 00:00:00 2001 From: Hieu Pham Date: Thu, 9 Jul 2026 07:14:45 -0700 Subject: [PATCH] Support Roaring-backed bitsets Add BitsetView support for dense, live Roaring, and frozen Roaring filters, including count updates, dense conversion, offsets, and out-id mapping. Preserve Roaring filters through search setup instead of rebuilding dense views, and add a typed Faiss fast path that expands Roaring filters into valid ID arrays for fp16, bf16, and int8 exhaustive distance paths. Wire roaring into Conan and CMake dependencies and cover live and frozen Roaring bitsets with unit tests. Signed-off-by: Hieu Pham --- CMakeLists.txt | 7 +- cmake/libs/libdiskann.cmake | 5 +- cmake/libs/libfaiss.cmake | 6 +- conanfile.py | 2 + include/knowhere/bitsetview.h | 101 +++++- src/index/index.cc | 13 +- tests/ut/test_distances.cc | 205 ++++++++++++ tests/ut/test_utils.cc | 64 ++++ .../cppcontrib/knowhere/utils/distances_if.h | 1 - .../knowhere/utils/distances_typed.cpp | 298 +++++++++--------- 10 files changed, 541 insertions(+), 161 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 99de33262..d42377e9a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -103,6 +103,8 @@ endif() find_package(xxHash REQUIRED) include_directories(${xxHash_INCLUDE_DIRS}) find_package(simde REQUIRED) +find_package(roaring REQUIRED) +include_directories(${roaring_INCLUDE_DIRS}) if(NOT WITH_LIGHT) find_package(opentelemetry-cpp REQUIRED) @@ -114,6 +116,7 @@ if(NOT WITH_LIGHT) endif() if(APPLE) + add_definitions(-D_DARWIN_C_SOURCE) # Prevent SDK sysroot include dir from appearing as -I/-isystem in compile # commands. Transitive find_package calls (CURL, ZLIB, etc.) resolve headers # to the SDK sysroot, and the explicit -I flag breaks C++ wrapper headers @@ -125,7 +128,7 @@ if(APPLE) endif() set(CMAKE_CXX_STANDARD 20) set(CMAKE_OSX_DEPLOYMENT_TARGET - "10.15" + "11.0" CACHE STRING "Minimum OS X deployment version" FORCE) if(OPENMP_FOUND) @@ -139,6 +142,7 @@ include(cmake/utils/compile_flags.cmake) find_package(milvus-common REQUIRED) include(cmake/libs/libhnsw.cmake) include(cmake/libs/libfaiss.cmake) +target_link_libraries(faiss PUBLIC roaring::roaring) if(WITH_COVERAGE) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fprofile-arcs -ftest-coverage") @@ -268,6 +272,7 @@ list(APPEND KNOWHERE_LINKER_LIBS fmt::fmt) list(APPEND KNOWHERE_LINKER_LIBS Folly::folly) list(APPEND KNOWHERE_LINKER_LIBS milvus-common::milvus-common) list(APPEND KNOWHERE_LINKER_LIBS simde::simde) +list(APPEND KNOWHERE_LINKER_LIBS roaring::roaring) add_library(knowhere SHARED ${KNOWHERE_SRCS}) add_dependencies(knowhere ${KNOWHERE_LINKER_LIBS}) diff --git a/cmake/libs/libdiskann.cmake b/cmake/libs/libdiskann.cmake index 9c6ce5eb3..150dba9c2 100644 --- a/cmake/libs/libdiskann.cmake +++ b/cmake/libs/libdiskann.cmake @@ -29,7 +29,7 @@ set(DISKANN_SOURCES find_package(folly REQUIRED) set(DISKANN_LINKER_LIBS PUBLIC ${AIO_LIBRARIES} ${DISKANN_BOOST_PROGRAM_OPTIONS_LIB} nlohmann_json::nlohmann_json - Folly::folly fmt::fmt-header-only prometheus-cpp::core prometheus-cpp::push glog::glog) + Folly::folly fmt::fmt-header-only prometheus-cpp::core prometheus-cpp::push glog::glog roaring::roaring) if (WITH_CUVS) list(APPEND DISKANN_LINKER_LIBS PRIVATE cuvs::cuvs) list(APPEND DISKANN_SOURCES thirdparty/DiskANN/src/diskann_gpu.cpp) @@ -49,7 +49,8 @@ target_link_libraries( fmt::fmt prometheus-cpp::core prometheus-cpp::push - glog::glog) + glog::glog + roaring::roaring) if(__X86_64) target_compile_options( diskann PRIVATE -fno-builtin-malloc -fno-builtin-calloc diff --git a/cmake/libs/libfaiss.cmake b/cmake/libs/libfaiss.cmake index 056b530e6..61944a2ff 100644 --- a/cmake/libs/libfaiss.cmake +++ b/cmake/libs/libfaiss.cmake @@ -384,7 +384,7 @@ if(__X86_64) -mavx2 -mfma -mf16c -mpopcnt>) target_compile_definitions(faiss_avx2 PRIVATE COMPILE_SIMD_AVX2) target_include_directories(faiss_avx2 PRIVATE ${Boost_INCLUDE_DIRS}) - target_link_libraries(faiss_avx2 PRIVATE milvus-common::milvus-common) + target_link_libraries(faiss_avx2 PRIVATE milvus-common::milvus-common roaring::roaring) add_library(faiss_avx512 OBJECT ${FAISS_AVX512_SRCS}) target_compile_options( faiss_avx512 @@ -400,7 +400,7 @@ if(__X86_64) -mpopcnt>) target_compile_definitions(faiss_avx512 PRIVATE COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512) target_include_directories(faiss_avx512 PRIVATE ${Boost_INCLUDE_DIRS}) - target_link_libraries(faiss_avx512 PRIVATE milvus-common::milvus-common) + target_link_libraries(faiss_avx512 PRIVATE milvus-common::milvus-common roaring::roaring) if(FAISS_ENABLE_AVX512_SPR) add_library(faiss_avx512_spr OBJECT ${FAISS_AVX512_SPR_SRCS}) @@ -424,7 +424,7 @@ if(__X86_64) target_compile_definitions(faiss_avx512_spr PRIVATE COMPILE_SIMD_AVX2 COMPILE_SIMD_AVX512 COMPILE_SIMD_AVX512_SPR) target_include_directories(faiss_avx512_spr PRIVATE ${Boost_INCLUDE_DIRS}) - target_link_libraries(faiss_avx512_spr PRIVATE milvus-common::milvus-common) + target_link_libraries(faiss_avx512_spr PRIVATE milvus-common::milvus-common roaring::roaring) endif() add_library(faiss STATIC ${FAISS_SRCS}) diff --git a/conanfile.py b/conanfile.py index 41e0edc6c..4fdb3a046 100644 --- a/conanfile.py +++ b/conanfile.py @@ -137,6 +137,7 @@ def requirements(self): self.requires("libcurl/8.10.1#a3113369c86086b0e84231844e7ed0a9", force=True, override=True) self.requires("simde/0.8.2#5e1edfd5cba92f25d79bf6ef4616b972") self.requires("xxhash/0.8.3#caa6d0af1b951c247922e38fbcebdbe6") + self.requires("roaring/3.0.0") if self.settings.os == "Linux": self.requires("openblas/0.3.30") if not self.options.with_light: @@ -283,6 +284,7 @@ def package_info(self): "milvus-common::milvus-common", "prometheus-cpp::core", "prometheus-cpp::push", + "roaring::roaring", ] self.cpp_info.components["libknowhere"].set_property( diff --git a/include/knowhere/bitsetview.h b/include/knowhere/bitsetview.h index ee6db3a47..315f12ec0 100644 --- a/include/knowhere/bitsetview.h +++ b/include/knowhere/bitsetview.h @@ -12,11 +12,15 @@ #ifndef BITSET_H #define BITSET_H +#include + #include #include +#include #include #include #include +#include namespace knowhere { class BitsetView { @@ -25,12 +29,38 @@ class BitsetView { ~BitsetView() = default; BitsetView(const uint8_t* data, size_t num_bits, size_t num_filtered_out_bits = 0, size_t id_offset = 0) - : bits_(data), num_bits_(num_bits), num_filtered_out_bits_(num_filtered_out_bits), id_offset_(id_offset) { + : kind_(Kind::Dense), + bits_(data), + num_bits_(num_bits), + num_filtered_out_bits_(num_filtered_out_bits), + id_offset_(id_offset) { + } + + BitsetView(const roaring_bitmap_t* bitmap, size_t num_bits, size_t num_filtered_out_bits = 0, size_t id_offset = 0) + : kind_(Kind::Roaring), + roaring_(bitmap), + num_bits_(num_bits), + num_filtered_out_bits_(num_filtered_out_bits), + id_offset_(id_offset) { } BitsetView(const std::nullptr_t) : BitsetView() { } + BitsetView(const std::nullptr_t, size_t num_bits, size_t num_filtered_out_bits = 0, size_t id_offset = 0) + : BitsetView(static_cast(nullptr), num_bits, num_filtered_out_bits, id_offset) { + } + + static BitsetView + FromFrozenRoaring(const void* data, size_t byte_size, size_t num_bits, size_t num_filtered_out_bits = 0, + size_t id_offset = 0) { + const auto* bitmap = roaring_bitmap_frozen_view(static_cast(data), byte_size); + BitsetView bitset(bitmap, num_bits, num_filtered_out_bits, id_offset); + bitset.owned_roaring_ = std::shared_ptr( + bitmap, [](const roaring_bitmap_t* p) { roaring_bitmap_free(const_cast(p)); }); + return bitset; + } + bool empty() const { return num_bits_ == 0; @@ -54,6 +84,15 @@ class BitsetView { return num_filtered_out_bits_; } + void + set_count(size_t num_filtered_out_bits) { + if (out_ids_ != nullptr) { + num_filtered_out_ids_ = num_filtered_out_bits; + return; + } + num_filtered_out_bits_ = num_filtered_out_bits; + } + size_t byte_size() const { return (num_bits_ + 8 - 1) >> 3; @@ -64,6 +103,42 @@ class BitsetView { return bits_; } + bool + is_dense() const { + return kind_ == Kind::Dense; + } + + bool + is_roaring() const { + return kind_ == Kind::Roaring; + } + + const roaring_bitmap_t* + roaring() const { + return roaring_; + } + + size_t + id_offset() const { + return id_offset_; + } + + bool + can_iterate_roaring_without_mapping() const { + return kind_ == Kind::Roaring && out_ids_ == nullptr; + } + + std::vector + ToDense() const { + std::vector dense(byte_size(), 0); + for (size_t i = 0; i < num_bits_; ++i) { + if (test(i)) { + dense[i >> 3] |= 0x1 << (i & 0x7); + } + } + return dense; + } + bool has_out_ids() const { return out_ids_ != nullptr; @@ -103,7 +178,13 @@ class BitsetView { out_id = out_ids_[out_id]; } // when index is larger than the max_offset, ignore it - return (out_id >= static_cast(num_bits_)) || (bits_[out_id >> 3] & (0x1 << (out_id & 0x7))); + if (out_id >= static_cast(num_bits_)) { + return true; + } + if (kind_ == Kind::Roaring) { + return roaring_bitmap_contains(roaring_, static_cast(out_id)); + } + return bits_[out_id >> 3] & (0x1 << (out_id & 0x7)); } // return the filtered ratio. if with id mapping, calculated by internal_ids rather than bits. float @@ -126,6 +207,9 @@ class BitsetView { } return count; } + if (kind_ == Kind::Roaring) { + return roaring_bitmap_get_cardinality(roaring_); + } // if without id mapping, use a better algorithm to calculate the number of filtered out bits. size_t ret = 0; auto len_uint8 = byte_size(); @@ -166,6 +250,14 @@ class BitsetView { } return num_internal_ids_; } + if (kind_ == Kind::Roaring) { + for (size_t i = 0; i < num_bits_; i++) { + if (!test(i)) { + return i; + } + } + return num_bits_; + } // if without id mapping, use a better algorithm to find the first valid index. size_t ret = 0; auto len_uint8 = byte_size(); @@ -211,7 +303,12 @@ class BitsetView { } private: + enum class Kind { Dense, Roaring }; + + Kind kind_ = Kind::Dense; const uint8_t* bits_ = nullptr; + const roaring_bitmap_t* roaring_ = nullptr; + std::shared_ptr owned_roaring_; size_t num_bits_ = 0; size_t num_filtered_out_bits_ = 0; diff --git a/src/index/index.cc b/src/index/index.cc index 70f615494..268a95e2e 100644 --- a/src/index/index.cc +++ b/src/index/index.cc @@ -154,7 +154,8 @@ Index::Search(const DataSetPtr dataset, const Json& json, const BitsetView& b if (bitset_.count() == 0) { // traverse bitset to get the filtered out num auto filtered_out_num = bitset_.get_filtered_out_num_(); - bitset = BitsetView(bitset_.data(), bitset_.size(), filtered_out_num); + bitset = bitset_; + bitset.set_count(filtered_out_num); } else { // if bitset has filtered out num, use it bitset = bitset_; @@ -212,7 +213,10 @@ Index::AnnIterator(const DataSetPtr dataset, const Json& json, const BitsetVi return expected>>::Err(Status::invalid_args, msg); } - const auto bitset = BitsetView(bitset_.data(), bitset_.size(), bitset_.get_filtered_out_num_()); + auto bitset = bitset_; + if (bitset.count() == 0) { + bitset.set_count(bitset.get_filtered_out_num_()); + } #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) // note that this time includes only the initial search phase of iterator. @@ -252,7 +256,10 @@ Index::RangeSearch(const DataSetPtr dataset, const Json& json, const BitsetVi return expected::Err(Status::invalid_args, msg); } - const auto bitset = BitsetView(bitset_.data(), bitset_.size(), bitset_.get_filtered_out_num_()); + auto bitset = bitset_; + if (bitset.count() == 0) { + bitset.set_count(bitset.get_filtered_out_num_()); + } #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) const BaseConfig& b_cfg = static_cast(*cfg); diff --git a/tests/ut/test_distances.cc b/tests/ut/test_distances.cc index f815b346e..0fe506e06 100644 --- a/tests/ut/test_distances.cc +++ b/tests/ut/test_distances.cc @@ -12,10 +12,77 @@ #include #include #include +#include +#include +#include +#include +#include #include +#include +#include +#include +#include "knowhere/operands.h" #include "simd/distances_ref.h" #include "simd/hook.h" + +namespace { + +template +void +CheckByIdxDistancePositions(ByIdxDistance by_idx_distance, DirectDistance direct_distance) { + constexpr size_t dim = 3; + constexpr size_t database_size = 32; + const std::array selected_ids = {5, 9, 17, 23, 31}; + const std::array query = {DataType(2.0f), DataType(-3.0f), DataType(5.0f)}; + std::vector database(database_size * dim); + for (size_t row = 0; row < database_size; ++row) { + database[row * dim] = DataType(static_cast(row % 11) - 5.0f); + database[row * dim + 1] = DataType(static_cast((row * 3) % 13) - 6.0f); + database[row * dim + 2] = DataType(static_cast((row * 5) % 17) - 8.0f); + } + + std::vector result_ids; + std::vector result_distances; + by_idx_distance( + query.data(), + database.data(), + selected_ids.data(), + dim, + selected_ids.size(), + [](size_t) -> std::optional { return true; }, + [&](float distance, int64_t id) { + result_ids.push_back(id); + result_distances.push_back(distance); + }); + + REQUIRE(result_ids == std::vector{0, 1, 2, 3, 4}); + REQUIRE(result_distances.size() == selected_ids.size()); + for (size_t i = 0; i < selected_ids.size(); ++i) { + const auto id = selected_ids[i]; + const float expected = direct_distance(query.data(), database.data() + id * dim, dim); + REQUIRE_THAT(result_distances[i], Catch::Matchers::WithinAbs(expected, 0.001f)); + } +} + +struct RecordingDistanceComputer : faiss::DistanceComputer { + void + set_query(const float*) override { + } + + float + operator()(faiss::idx_t id) override { + return static_cast(id); + } + + float + symmetric_dis(faiss::idx_t, faiss::idx_t) override { + return 0.0f; + } +}; + +} // namespace + TEST_CASE("Test Distance Compute", "[distance]") { std::mt19937 rng; std::uniform_int_distribution<> distrib(1, 100000); @@ -122,3 +189,141 @@ TEST_CASE("Test Distance Compute", "[distance]") { } } } + +TEST_CASE("Test typed by-idx distance callbacks use positions", "[distance][typed][by_idx]") { + using namespace faiss::cppcontrib::knowhere; + + SECTION("INT8 inner product") { + CheckByIdxDistancePositions<::knowhere::int8>( + [](auto... args) { int8_vec_inner_products_ny_by_idx_if(args...); }, + int8_vec_inner_product); + } + + SECTION("INT8 L2") { + CheckByIdxDistancePositions<::knowhere::int8>( + [](auto... args) { int8_vec_L2sqr_ny_by_idx_if(args...); }, + int8_vec_L2sqr); + } + + SECTION("FP16 inner product") { + CheckByIdxDistancePositions<::knowhere::fp16>( + [](auto... args) { fp16_vec_inner_products_ny_by_idx_if(args...); }, + fp16_vec_inner_product); + } + + SECTION("FP16 L2") { + CheckByIdxDistancePositions<::knowhere::fp16>( + [](auto... args) { fp16_vec_L2sqr_ny_by_idx_if(args...); }, + fp16_vec_L2sqr); + } + + SECTION("BF16 inner product") { + CheckByIdxDistancePositions<::knowhere::bf16>( + [](auto... args) { bf16_vec_inner_products_ny_by_idx_if(args...); }, + bf16_vec_inner_product); + } + + SECTION("BF16 L2") { + CheckByIdxDistancePositions<::knowhere::bf16>( + [](auto... args) { bf16_vec_L2sqr_ny_by_idx_if(args...); }, + bf16_vec_L2sqr); + } +} + +TEST_CASE("Test generic by-idx distance callbacks use positions", "[distance][by_idx]") { + const std::array ids = {5, 9, 17}; + RecordingDistanceComputer distance_computer; + std::vector callback_positions; + std::vector distances; + + faiss::cppcontrib::knowhere::distance_compute_by_idx_if( + ids.data(), + ids.size(), + &distance_computer, + [](size_t) -> std::optional { return true; }, + [&](float distance, faiss::idx_t position) { + distances.push_back(distance); + callback_positions.push_back(position); + }); + + REQUIRE(callback_positions == std::vector{0, 1, 2}); + REQUIRE(distances == std::vector{5.0f, 9.0f, 17.0f}); +} + +TEST_CASE("Test typed selector results use database IDs", "[distance][typed][selector]") { + using namespace faiss::cppcontrib::knowhere; + constexpr size_t dim = 3; + constexpr size_t database_size = 18; + const std::array selected_ids = {5, 9, 17}; + const std::array<::knowhere::int8, dim> query = {2, -3, 5}; + std::vector<::knowhere::int8> database(database_size * dim); + for (size_t row = 0; row < database_size; ++row) { + database[row * dim] = static_cast<::knowhere::int8>(row % 7 + 1); + database[row * dim + 1] = static_cast<::knowhere::int8>(row % 5 + 2); + database[row * dim + 2] = static_cast<::knowhere::int8>(row % 3 + 3); + } + faiss::IDSelectorArray selector(selected_ids.size(), selected_ids.data()); + + auto check_results = [&](const std::array& distances, + const std::array& result_ids, + auto expected_distance) { + auto sorted_result_ids = result_ids; + std::sort(sorted_result_ids.begin(), sorted_result_ids.end()); + REQUIRE(sorted_result_ids == selected_ids); + for (size_t i = 0; i < result_ids.size(); ++i) { + REQUIRE_THAT(distances[i], Catch::Matchers::WithinAbs(expected_distance(result_ids[i]), 0.001f)); + } + }; + + SECTION("inner product") { + std::array distances; + std::array result_ids; + knn_inner_product_typed( + query.data(), database.data(), dim, 1, database_size, 3, distances.data(), result_ids.data(), &selector); + check_results(distances, result_ids, [&](int64_t id) { + return int8_vec_inner_product(query.data(), database.data() + id * dim, dim); + }); + } + + SECTION("L2") { + std::array distances; + std::array result_ids; + knn_L2sqr_typed(query.data(), + database.data(), + dim, + 1, + database_size, + 3, + distances.data(), + result_ids.data(), + nullptr, + &selector); + check_results(distances, result_ids, [&](int64_t id) { + return int8_vec_L2sqr(query.data(), database.data() + id * dim, dim); + }); + } + + SECTION("cosine") { + std::vector y_inv_norms(database_size); + for (size_t id = 0; id < database_size; ++id) { + y_inv_norms[id] = 1.0f / sqrtf(int8_vec_norm_L2sqr(database.data() + id * dim, dim)); + } + const float x_inv_norm = 1.0f / sqrtf(int8_vec_norm_L2sqr(query.data(), dim)); + std::array distances; + std::array result_ids; + knn_cosine_typed(query.data(), + database.data(), + y_inv_norms.data(), + dim, + 1, + database_size, + 3, + distances.data(), + result_ids.data(), + &selector); + check_results(distances, result_ids, [&](int64_t id) { + return int8_vec_inner_product(query.data(), database.data() + id * dim, dim) * x_inv_norm * + y_inv_norms[id]; + }); + } +} diff --git a/tests/ut/test_utils.cc b/tests/ut/test_utils.cc index e14053ed3..858ef1e76 100644 --- a/tests/ut/test_utils.cc +++ b/tests/ut/test_utils.cc @@ -21,6 +21,7 @@ #include "knowhere/heap.h" #include "knowhere/utils.h" #include "knowhere/version.h" +#include "roaring/roaring.h" #include "utils.h" namespace { @@ -113,6 +114,69 @@ TEST_CASE("Test Bitset Generation", "[utils]") { } } +TEST_CASE("Test BitsetView Roaring", "[utils]") { + std::vector dense(2, 0); + for (auto id : {1, 3, 8, 13}) { + dense[id >> 3] |= 0x1 << (id & 0x7); + } + + auto* roaring = roaring_bitmap_create(); + for (auto id : {1, 3, 8, 13}) { + roaring_bitmap_add(roaring, id); + } + + knowhere::BitsetView dense_view(dense.data(), 16); + knowhere::BitsetView roaring_view(roaring, 16); + REQUIRE(roaring_view.is_roaring()); + REQUIRE(roaring_view.count() == 0); + REQUIRE(roaring_view.get_filtered_out_num_() == dense_view.get_filtered_out_num_()); + REQUIRE(roaring_view.filter_ratio() == dense_view.filter_ratio()); + REQUIRE(roaring_view.get_first_valid_index() == dense_view.get_first_valid_index()); + for (size_t i = 0; i < 16; ++i) { + REQUIRE(roaring_view.test(i) == dense_view.test(i)); + } + REQUIRE(roaring_view.ToDense() == dense); + + roaring_view.set_id_offset(2); + dense_view.set_id_offset(2); + for (size_t i = 0; i < 14; ++i) { + REQUIRE(roaring_view.test(i) == dense_view.test(i)); + } + + std::vector out_ids{0, 1, 3, 7, 8, 13}; + knowhere::BitsetView roaring_out_ids_view(roaring, 16); + knowhere::BitsetView dense_out_ids_view(dense.data(), 16); + roaring_out_ids_view.set_out_ids(out_ids.data(), out_ids.size()); + dense_out_ids_view.set_out_ids(out_ids.data(), out_ids.size()); + REQUIRE(roaring_out_ids_view.count() == dense_out_ids_view.count()); + for (size_t i = 0; i < out_ids.size(); ++i) { + REQUIRE(roaring_out_ids_view.test(i) == dense_out_ids_view.test(i)); + } + + roaring_bitmap_free(roaring); +} + +TEST_CASE("Test BitsetView Frozen Roaring", "[utils]") { + auto* roaring = roaring_bitmap_create(); + for (auto id : {0, 2, 4, 10}) { + roaring_bitmap_add(roaring, id); + } + const auto frozen_size = roaring_bitmap_frozen_size_in_bytes(roaring); + std::vector frozen(frozen_size); + roaring_bitmap_frozen_serialize(roaring, frozen.data()); + + auto frozen_view = knowhere::BitsetView::FromFrozenRoaring(frozen.data(), frozen.size(), 12); + REQUIRE(frozen_view.is_roaring()); + REQUIRE(frozen_view.get_filtered_out_num_() == 4); + REQUIRE(frozen_view.get_first_valid_index() == 1); + REQUIRE(frozen_view.test(0)); + REQUIRE(!frozen_view.test(1)); + REQUIRE(frozen_view.test(10)); + REQUIRE(frozen_view.test(12)); + + roaring_bitmap_free(roaring); +} + namespace { constexpr size_t kHeapSize = 10; constexpr size_t kElementCount = 10000; diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_if.h b/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_if.h index 082128c8f..c5e580eef 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_if.h +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_if.h @@ -1270,4 +1270,3 @@ void int8_vec_L2sqr_ny_by_idx_if( } } } //namespace faiss - diff --git a/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_typed.cpp b/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_typed.cpp index f300de47e..af2a1a6a7 100644 --- a/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_typed.cpp +++ b/thirdparty/faiss/faiss/cppcontrib/knowhere/utils/distances_typed.cpp @@ -14,6 +14,8 @@ #include #include +#include +#include #include #include @@ -30,6 +32,61 @@ namespace faiss::cppcontrib::knowhere { namespace { +bool make_roaring_valid_ids( + const ::knowhere::BitsetView& bitset_view, + size_t ny, + std::vector& ids) { + ids.clear(); + if (!bitset_view.can_iterate_roaring_without_mapping()) { + return false; + } + + const uint64_t range_start = bitset_view.id_offset(); + const uint64_t range_end = range_start + ny; + std::unique_ptr range( + roaring_bitmap_from_range(range_start, range_end, 1), + roaring_bitmap_free); + if (range == nullptr) { + return false; + } + std::unique_ptr selected( + roaring_bitmap_andnot(range.get(), bitset_view.roaring()), + roaring_bitmap_free); + if (selected == nullptr) { + return false; + } + + const auto nselected = roaring_bitmap_get_cardinality(selected.get()); + ids.resize(nselected); + std::vector selected_ids(nselected); + roaring_bitmap_to_uint32_array(selected.get(), selected_ids.data()); + for (size_t i = 0; i < nselected; ++i) { + ids[i] = static_cast(selected_ids[i] - range_start); + } + return true; +} + +template +void dispatch_typed_with_roaring_fast_path( + const IDSelector* sel, + size_t ny, + Fn&& fn) { + if (const auto* sel_bs = + dynamic_cast(sel)) { + std::vector roaring_ids; + if (make_roaring_valid_ids(sel_bs->bitset_view, ny, roaring_ids)) { + IDSelectorArray sel_array(roaring_ids.size(), roaring_ids.data()); + fn(sel_array); + } else { + fn(*sel_bs); + } + } else if (sel == nullptr) { + fn(IDSelectorAll()); + } else { + fn(*sel); + } +} + template void exhaustive_inner_product_impl_typed( const DataType* __restrict x, @@ -52,17 +109,20 @@ void exhaustive_inner_product_impl_typed( resi.add_result(ip, j); }; if constexpr (std::is_same_v) { - // todo: need more tests about this branch auto filter = [](const size_t j) { return true; }; + auto apply_by_idx = [&resi, &selector](const float distance, + const idx_t position) { + resi.add_result(distance, selector.ids[position]); + }; if constexpr (std::is_same_v) { fp16_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector.n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { bf16_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector.n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { int8_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector.n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } } else { // the lambda that filters acceptable elements. @@ -103,17 +163,20 @@ void exhaustive_L2sqr_seq_impl_typed( resi.add_result(ip, j); }; if constexpr (std::is_same_v) { - // todo: need more tests about this branch auto filter = [](const size_t j) { return true; }; + auto apply_by_idx = [&resi, &selector](const float distance, + const idx_t position) { + resi.add_result(distance, selector.ids[position]); + }; if constexpr (std::is_same_v) { fp16_vec_L2sqr_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { bf16_vec_L2sqr_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { int8_vec_L2sqr_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } } else { // the lambda that filters acceptable elements. @@ -170,17 +233,34 @@ void exhaustive_cosine_seq_impl_typed( resi.begin(i); // the lambda that applies a filtered element if constexpr (std::is_same_v) { - // todo: need more tests about this branch auto filter = [](const size_t j) { return true; }; + auto apply_by_idx = [&resi, + &selector, + x_inv_norm, + y, + y_inv_norms, + d, + norm_computer](const float ip, + const idx_t position) { + const auto id = selector.ids[position]; + const float y_inv_norm = (y_inv_norms != nullptr) + ? y_inv_norms[id] + : 1.0f / + std::max( + sqrtf(norm_computer(y + id * d, d)), + FLT_MIN); + + resi.add_result(ip * (x_inv_norm * y_inv_norm), id); + }; if constexpr (std::is_same_v) { fp16_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { bf16_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } else if constexpr (std::is_same_v) { int8_vec_inner_products_ny_by_idx_if( - x_i, y, selector.ids, d, selector->n, filter, apply); + x_i, y, selector.ids, d, selector.n, filter, apply_by_idx); } } else { // the lambda that filters acceptable elements. @@ -221,26 +301,18 @@ void knn_inner_product_typed( } if (k < distance_compute_min_k_reservoir) { HeapBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_inner_product_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_inner_product_impl_typed( + x, y, d, nx, ny, res, selector); + }); } else { ReservoirBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_inner_product_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_inner_product_impl_typed( + x, y, d, nx, ny, res, selector); + }); } if (imin != 0) { @@ -266,15 +338,9 @@ void all_inner_product_typed( const std::pair sentinel{-1, CMax::neutral()}; std::vector> pairs(n, sentinel); CollectAllResultHandler> res(nx, ny, pairs.data()); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_inner_product_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, selector); + }); pairs_to_distids(pairs.data(), output, n); } @@ -288,15 +354,9 @@ void all_inner_product_distances_typed( float* output, const IDSelector* sel) { CollectAllDistancesHandler> res(nx, ny, output); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_inner_product_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_inner_product_impl_typed(x, y, d, nx, ny, res, selector); + }); return; } @@ -322,26 +382,18 @@ void knn_L2sqr_typed( } if (k < distance_compute_min_k_reservoir) { HeapBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_L2sqr_seq_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_L2sqr_seq_impl_typed( + x, y, d, nx, ny, res, selector); + }); } else { ReservoirBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_L2sqr_seq_impl_typed( - x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_L2sqr_seq_impl_typed( + x, y, d, nx, ny, res, selector); + }); } if (imin != 0) { for (size_t i = 0; i < nx * k; i++) { @@ -367,14 +419,9 @@ void all_L2sqr_typed( const std::pair sentinel{-1, CMax::neutral()}; std::vector> pairs(n, sentinel); CollectAllResultHandler> res(nx, ny, pairs.data()); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, selector); + }); pairs_to_distids(pairs.data(), output, n); } @@ -389,14 +436,9 @@ void all_L2sqr_distances_typed( const float* y_norms, const IDSelector* sel) { CollectAllDistancesHandler> res(nx, ny, output); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, res, selector); + }); return; } @@ -422,30 +464,18 @@ void knn_cosine_typed( } if (k < distance_compute_min_k_reservoir) { HeapBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_cosine_seq_impl_typed( + x, y, y_inv_norms, d, nx, ny, res, selector); + }); } else { ReservoirBlockResultHandler> res(nx, vals, ids, k); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel); - } + dispatch_typed_with_roaring_fast_path( + sel, ny, [&](const auto& selector) { + exhaustive_cosine_seq_impl_typed( + x, y, y_inv_norms, d, nx, ny, res, selector); + }); } if (imin != 0) { for (size_t i = 0; i < nx * k; i++) { @@ -471,16 +501,10 @@ void all_cosine_typed( const std::pair sentinel{-1, CMax::neutral()}; std::vector> pairs(n, sentinel); CollectAllResultHandler> res(nx, ny, pairs.data()); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_cosine_seq_impl_typed(x, y, y_inv_norms, d, nx, ny, res, *sel); - } + x, y, y_inv_norms, d, nx, ny, res, selector); + }); pairs_to_distids(pairs.data(), output, n); } @@ -495,16 +519,10 @@ void all_cosine_distances_typed( float* distances, const IDSelector* sel) { CollectAllDistancesHandler> res(nx, ny, distances); - if (const auto* sel_bs = - dynamic_cast(sel)) { + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, *sel_bs); - } else if (sel == nullptr) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, res, IDSelectorAll()); - } else { - exhaustive_cosine_seq_impl_typed(x, y, y_inv_norms, d, nx, ny, res, *sel); - } + x, y, y_inv_norms, d, nx, ny, res, selector); + }); return; } @@ -523,14 +541,9 @@ void range_search_L2sqr_typed( faiss::RangeSearchResult* res, const IDSelector* sel) { RangeSearchBlockResultHandler> resh(res, radius); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, resh, *sel_bs); - } else if (sel == nullptr) { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, resh, IDSelectorAll()); - } else { - exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, resh, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_L2sqr_seq_impl_typed(x, y, d, nx, ny, resh, selector); + }); return; } @@ -545,15 +558,9 @@ void range_search_inner_product_typed( faiss::RangeSearchResult* res, const IDSelector* sel) { RangeSearchBlockResultHandler> resh(res, radius); - if (const auto* sel_bs = - dynamic_cast(sel)) { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, resh, *sel_bs); - } else if (sel == nullptr) { - exhaustive_inner_product_impl_typed( - x, y, d, nx, ny, resh, IDSelectorAll()); - } else { - exhaustive_inner_product_impl_typed(x, y, d, nx, ny, resh, *sel); - } + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { + exhaustive_inner_product_impl_typed(x, y, d, nx, ny, resh, selector); + }); return; } @@ -570,16 +577,10 @@ void range_search_cosine_typed( faiss::RangeSearchResult* res, const IDSelector* sel) { RangeSearchBlockResultHandler> resh(res, radius); - if (const auto* sel_bs = - dynamic_cast(sel)) { + dispatch_typed_with_roaring_fast_path(sel, ny, [&](const auto& selector) { exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, resh, *sel_bs); - } else if (sel == nullptr) { - exhaustive_cosine_seq_impl_typed( - x, y, y_inv_norms, d, nx, ny, resh, IDSelectorAll()); - } else { - exhaustive_cosine_seq_impl_typed(x, y, y_inv_norms, d, nx, ny, resh, *sel); - } + x, y, y_inv_norms, d, nx, ny, resh, selector); + }); return; } @@ -962,4 +963,3 @@ template void faiss::cppcontrib::knowhere::range_search_cosine_typed<::knowhere: const faiss::IDSelector*); } -