enhance: knowhere support metric mhjaccard and index minhash lsh for minhash vector (#1203)

Signed-off-by: cqy123456 <qianya.cheng@zilliz.com>

Author: cqy123456

cmake/libs/libfaiss.cmake

@@ -50,6 +50,7 @@ if(__X86_64)
     ${UTILS_SRC} $<TARGET_OBJECTS:utils_sse> $<TARGET_OBJECTS:utils_avx>
     $<TARGET_OBJECTS:utils_avx512> $<TARGET_OBJECTS:utils_avx512icx>)
   target_link_libraries(knowhere_utils PUBLIC glog::glog)
+  target_link_libraries(knowhere_utils PUBLIC xxHash::xxhash)
 endif()
 
 if(__AARCH64)
@@ -99,19 +100,22 @@ if(__AARCH64)
   endif()
 
   target_link_libraries(knowhere_utils PUBLIC glog::glog)
+  target_link_libraries(knowhere_utils PUBLIC xxHash::xxhash)
 endif()
 
 if(__RISCV64)
   set(UTILS_SRC src/simd/hook.cc src/simd/distances_ref.cc)
   add_library(knowhere_utils STATIC ${UTILS_SRC})
   target_link_libraries(knowhere_utils PUBLIC glog::glog)
+  target_link_libraries(knowhere_utils PUBLIC xxHash::xxhash)
 endif()
 
 # ToDo: Add distances_vsx.cc for powerpc64 SIMD acceleration
 if(__PPC64)
   set(UTILS_SRC src/simd/hook.cc src/simd/distances_ref.cc src/simd/distances_powerpc.cc)
   add_library(knowhere_utils STATIC ${UTILS_SRC})
   target_link_libraries(knowhere_utils PUBLIC glog::glog)
+  target_link_libraries(knowhere_utils PUBLIC xxHash::xxhash)
 endif()
 
 
@@ -131,6 +135,9 @@ else()
   find_package(BLAS REQUIRED)
 endif()
 
+find_package(xxHash REQUIRED)
+include_directories(${xxHash_INCLUDE_DIRS})
+
 if(__X86_64)
   list(REMOVE_ITEM FAISS_SRCS ${FAISS_AVX2_SRCS})
 

conanfile.py

@@ -106,6 +106,7 @@ def requirements(self):
         self.requires("folly/2023.10.30.09@milvus/dev")
         self.requires("libcurl/8.2.1")
         self.requires("simde/0.8.2")
+        self.requires("xxhash/0.8.2")
         if self.settings.os == "Android":
             self.requires("openblas/0.3.27")
         if not self.options.with_light:

include/knowhere/comp/bloomfilter.h

@@ -0,0 +1,121 @@
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+#ifndef KNOWHERE_KNOWHERE_H
+#define KNOWHERE_KNOWHERE_H
+#include <cmath>
+#include <cstring>
+#include <fstream>
+#include <functional>
+#include <vector>
+
+#include "io/memory_io.h"
+#include "knowhere/bitsetview.h"
+#include "knowhere/utils.h"
+namespace knowhere {
+template <typename T>
+class BloomFilter {
+ public:
+    BloomFilter() {
+    }
+    BloomFilter(size_t expected_elements, double false_positive_prob) : n(expected_elements), p(false_positive_prob) {
+        m = static_cast<size_t>(-(n * log(p)) / (log(2) * log(2)));
+        k = static_cast<int>(m / n * log(2));
+        m = std::max<size_t>(m, 1);
+        k = std::max(k, 1);
+        bits.resize(m, false);
+    }
+
+    void
+    add(const T& element) {
+        size_t glb_hash = hash((const char*)&element, sizeof(element), 0);
+        for (int i = 0; i < k; ++i) {
+            size_t pos = (glb_hash + i) % m;
+            bits[pos] = true;
+        }
+    }
+
+    bool
+    contains(const T& element) const {
+        size_t glb_hash = hash((const char*)&element, sizeof(element), 0);
+        for (int i = 0; i < k; ++i) {
+            size_t pos = (glb_hash + i) % m;
+            if (!bits[pos])
+                return false;
+        }
+        return true;
+    }
+
+    void
+    save(MemoryIOWriter& writer) const {
+        writeBinaryPOD(writer, m);
+        writeBinaryPOD(writer, k);
+        writeBinaryPOD(writer, n);
+        writeBinaryPOD(writer, p);
+        auto bytes_num = (m + 8 - 1) / 8;
+        std::vector<char> buffer(bytes_num, 0);
+        for (size_t i = 0; i < m; ++i) {
+            if (bits[i]) {
+                buffer[i / 8] |= (1 << (i % 8));
+            }
+        }
+        writer.write(buffer.data(), buffer.size());
+    }
+
+    void
+    load(MemoryIOReader& reader) {
+        readBinaryPOD(reader, m);
+        readBinaryPOD(reader, k);
+        readBinaryPOD(reader, n);
+        readBinaryPOD(reader, p);
+        bits.clear();
+        bits.resize(m);
+        auto bytes_num = (m + 8 - 1) / 8;
+        std::vector<char> buffer(bytes_num);
+        reader.read(buffer.data(), bytes_num);
+        for (size_t i = 0; i < m; ++i) {
+            bool bit = (buffer[i / 8] >> (i % 8)) & 1;
+            bits.push_back(bit);
+        }
+    }
+    size_t
+    size() const {
+        return n;
+    }
+    double
+    false_positive_rate() const {
+        return p;
+    }
+    size_t
+    memory_usage() const {
+        return m / 8;
+    }
+
+ private:
+    static constexpr size_t multiplier = 31;
+    std::vector<bool> bits;
+    size_t m = 0;
+    int k = 0;
+    double p = 0;
+    size_t n = 0;
+
+    size_t
+    hash(const char* data, size_t length, size_t bucket_i) const {
+        if (data == nullptr) {
+            throw std::runtime_error("can't hash null data.");
+        }
+        size_t result = 0;
+        for (size_t i = 0; i < length; ++i) {
+            result = (result * multiplier) + static_cast<size_t>(data[i]);
+        }
+        return (result + bucket_i) % m;
+    }
+};
+}  // namespace knowhere
+#endif

include/knowhere/comp/index_param.h

@@ -56,6 +56,7 @@ constexpr const char* INDEX_HNSW_PQ = "HNSW_PQ";
 constexpr const char* INDEX_HNSW_PRQ = "HNSW_PRQ";
 
 constexpr const char* INDEX_DISKANN = "DISKANN";
+constexpr const char* INDEX_MINHASH_INDEX = "MINHASH_LSH";
 
 constexpr const char* INDEX_SPARSE_INVERTED_INDEX = "SPARSE_INVERTED_INDEX";
 constexpr const char* INDEX_SPARSE_WAND = "SPARSE_WAND";
@@ -187,6 +188,16 @@ constexpr const char* DROP_RATIO_SEARCH = "drop_ratio_search";
 // RaBitQ Params
 constexpr const char* RABITQ_QUERY_BITS = "rbq_bits_query";
 
+// minhash Params
+constexpr const char* ALIGNED_BLOCK_SIZE = "aligned_block_size";
+constexpr const char* BAND = "band";
+constexpr const char* SHARED_BLOOM_FILTER = "shared_bloom_filter";
+constexpr const char* BLOOM_FALSE_POSITIVE_RPOB = "bloom_false_positive_prob";
+constexpr const char* HASH_CODE_IN_MEM = "hash_code_in_mem";
+constexpr const char* SEARCH_WITH_JACCARD = "search_with_jaccard";
+constexpr const char* MH_ELEMENT_BIT_WIDTH = "mh_element_bit_width";
+constexpr const char* MH_LSH_REFINE_K = "refine_k";
+constexpr const char* BATCH_SEARCH = "batch_search";
 }  // namespace indexparam
 
 using MetricType = std::string;
@@ -197,6 +208,7 @@ constexpr const char* L2 = "L2";
 constexpr const char* COSINE = "COSINE";
 constexpr const char* HAMMING = "HAMMING";
 constexpr const char* JACCARD = "JACCARD";
+constexpr const char* MHJACCARD = "MHJACCARD";
 constexpr const char* SUBSTRUCTURE = "SUBSTRUCTURE";
 constexpr const char* SUPERSTRUCTURE = "SUPERSTRUCTURE";
 constexpr const char* BM25 = "BM25";

include/knowhere/config.h

@@ -633,6 +633,12 @@ class BaseConfig : public Config {
      */
     CFG_INT refine_type;
     CFG_BOOL refine_with_quant;
+    /*
+     * band is a special parameters of BF search and MinHash index node train.
+     */
+    CFG_INT band;
+    CFG_BOOL search_with_jaccard;
+    CFG_INT mh_element_bit_width;
     KNOHWERE_DECLARE_CONFIG(BaseConfig) {
         KNOWHERE_CONFIG_DECLARE_FIELD(dim).allow_empty_without_default().description("vector dim").for_train();
         KNOWHERE_CONFIG_DECLARE_FIELD(metric_type)
@@ -781,6 +787,17 @@ class BaseConfig : public Config {
             .for_search()
             .for_range_search()
             .for_iterator();
+        KNOWHERE_CONFIG_DECLARE_FIELD(band).description("param of MinHashLSH").set_default(1).for_train().for_search();
+        KNOWHERE_CONFIG_DECLARE_FIELD(mh_element_bit_width)
+            .description("sizeof(hash code), the hash element should be aligned on 8 bits")
+            .set_default(8)
+            .set_range(8, 256)
+            .for_train()
+            .for_search();
+        KNOWHERE_CONFIG_DECLARE_FIELD(search_with_jaccard)
+            .description("return the jaccard distance of minhash vector search or minhashlsh hit flag.")
+            .set_default(false)
+            .for_search();
     }
 };
 }  // namespace knowhere

include/knowhere/index/index_factory.h

@@ -143,6 +143,10 @@ class IndexFactory {
 #define KNOWHERE_MOCK_REGISTER_DENSE_INT_GLOBAL(name, index_node, features, ...) \
     KNOWHERE_MOCK_REGISTER_GLOBAL(name, index_node, int8, (features | knowhere::feature::INT8), ##__VA_ARGS__);
 
+// register vector index supporting binary data types
+#define KNOWHERE_MOCK_REGISTER_DENSE_BINARY_ALL_GLOBAL(name, index_node, features, ...) \
+    KNOWHERE_SIMPLE_REGISTER_GLOBAL(name, index_node, bin1, (features | knowhere::feature::BINARY), ##__VA_ARGS__);
+
 // register vector index supporting ALL_DENSE_FLOAT_TYPE(float32, bf16, fp16) data types, but mocked bf16 and fp16
 #define KNOWHERE_MOCK_REGISTER_DENSE_FLOAT_ALL_GLOBAL(name, index_node, features, ...)                          \
     KNOWHERE_MOCK_REGISTER_GLOBAL(name, index_node, bf16, (features | knowhere::feature::BF16), ##__VA_ARGS__); \

include/knowhere/index/index_table.h

@@ -98,6 +98,8 @@ static std::set<std::pair<std::string, VecType>> legal_knowhere_index = {
     // sparse index
     {IndexEnum::INDEX_SPARSE_INVERTED_INDEX, VecType::VECTOR_SPARSE_FLOAT},
     {IndexEnum::INDEX_SPARSE_WAND, VecType::VECTOR_SPARSE_FLOAT},
+    //  minhash index
+    {IndexEnum::INDEX_MINHASH_INDEX, VecType::VECTOR_BINARY},
 };
 
 static std::set<std::string> legal_support_mmap_knowhere_index = {

include/knowhere/utils.h

@@ -14,6 +14,7 @@
 #include <strings.h>
 
 #include <algorithm>
+#include <fstream>
 #include <vector>
 
 #include "knowhere/binaryset.h"
@@ -226,6 +227,25 @@ readBinaryPOD(R& in, T& podRef) {
     in.read((char*)&podRef, sizeof(T));
 }
 
+inline void
+load_binary_vec(const std::string& bin_file, std::unique_ptr<char[]>& data, size_t& npts, size_t& dim) {
+    std::ifstream file(bin_file, std::ios::binary);
+    if (!file.is_open()) {
+        throw std::runtime_error("fail to open file: " + bin_file);
+    }
+    uint32_t n, d;
+    file.read(reinterpret_cast<char*>(&n), sizeof(uint32_t));
+    file.read(reinterpret_cast<char*>(&d), sizeof(uint32_t));
+    npts = n;
+    dim = d;
+    if (dim % 8 != 0) {
+        throw std::runtime_error("fail to load binary vector base file, dim % 8 != 0 ");
+    }
+    uint64_t total_size = dim * npts / 8;
+    data = std::make_unique<char[]>(total_size);
+    file.read(reinterpret_cast<char*>(data.get()), total_size);
+}
+
 // taken from
 // https://github.com/Microsoft/BLAS-on-flash/blob/master/include/utils.h
 // round up X to the nearest multiple of Y

src/common/comp/brute_force.cc

@@ -18,6 +18,7 @@
 #include "faiss/utils/binary_distances.h"
 #include "faiss/utils/distances.h"
 #include "faiss/utils/distances_typed.h"
+#include "index/minhash/minhash_util.h"
 #include "knowhere/bitsetview_idselector.h"
 #include "knowhere/comp/thread_pool.h"
 #include "knowhere/config.h"
@@ -151,6 +152,14 @@ BruteForce::Search(const DataSetPtr base_dataset, const DataSetPtr query_dataset
     auto labels = std::make_unique<int64_t[]>(nq * topk);
     auto distances = std::make_unique<float[]>(nq * topk);
     std::unique_ptr<float[]> norms = is_cosine ? GetVecNorms<DataType>(base_dataset) : nullptr;
+    // some check for minhash metric
+    if (faiss_metric_type == faiss::METRIC_MinHash_Jaccard) {
+        auto mh_valid_stat =
+            MinhashConfigCheck(dim, datatype_v<DataType>, PARAM_TYPE::SEARCH | PARAM_TYPE::TRAIN, &cfg, &bitset);
+        if (mh_valid_stat != Status::success) {
+            return expected<DataSetPtr>::Err(mh_valid_stat, "MinhashConfigCheck() failed, please check the config.");
+        }
+    }
     auto pool = ThreadPool::GetGlobalSearchThreadPool();
     std::vector<folly::Future<Status>> futs;
     futs.reserve(nq);
@@ -213,6 +222,23 @@ BruteForce::Search(const DataSetPtr base_dataset, const DataSetPtr query_dataset
                     binary_knn_hc(faiss::METRIC_Jaccard, &res, cur_query, (const uint8_t*)xb, nb, dim / 8, id_selector);
                     break;
                 }
+                case faiss::METRIC_MinHash_Jaccard: {
+                    size_t band = cfg.band.value();
+                    bool search_with_jaccard = cfg.search_with_jaccard.value();
+                    if (search_with_jaccard) {
+                        size_t hash_element_size = cfg.mh_element_bit_width.value() / 8;  // in bytes
+                        size_t hash_element_length = dim / (hash_element_size * 8);
+                        auto cur_query = (const char*)xq + (dim / 8) * index;
+                        minhash_jaccard_knn_ny(cur_query, (const char*)xb, hash_element_length, hash_element_size, nb,
+                                               topk, bitset, cur_distances, cur_labels);
+                    } else {
+                        size_t u8_dim = dim / 8;
+                        auto cur_query = (const char*)xq + u8_dim * index;
+                        minhash_lsh_hit_ny(cur_query, (const char*)xb, u8_dim, band, nb, topk, bitset, cur_distances,
+                                           cur_labels);
+                    }
+                    break;
+                }
                 case faiss::METRIC_Hamming: {
                     auto cur_query = (const uint8_t*)xq + (dim / 8) * index;
                     std::vector<int32_t> int_distances(topk);
@@ -306,6 +332,14 @@ BruteForce::SearchWithBuf(const DataSetPtr base_dataset, const DataSetPtr query_
     int topk = cfg.k.value();
     auto labels = ids;
     auto distances = dis;
+    // some check for minhash metric
+    if (faiss_metric_type == faiss::METRIC_MinHash_Jaccard) {
+        auto mh_valid_stat =
+            MinhashConfigCheck(dim, datatype_v<DataType>, PARAM_TYPE::SEARCH | PARAM_TYPE::TRAIN, &cfg, &bitset);
+        if (mh_valid_stat != Status::success) {
+            return mh_valid_stat;
+        }
+    }
 
     std::unique_ptr<float[]> norms = is_cosine ? GetVecNorms<DataType>(base_dataset) : nullptr;
     auto pool = ThreadPool::GetGlobalSearchThreadPool();
@@ -363,6 +397,23 @@ BruteForce::SearchWithBuf(const DataSetPtr base_dataset, const DataSetPtr query_
                     }
                     break;
                 }
+                case faiss::METRIC_MinHash_Jaccard: {
+                    size_t band = cfg.band.value();
+                    bool search_with_jaccard = cfg.search_with_jaccard.value();
+                    if (search_with_jaccard) {
+                        size_t hash_element_size = cfg.mh_element_bit_width.value() / 8;  // in bytes
+                        size_t hash_element_length = dim / (hash_element_size * 8);
+                        auto cur_query = (const char*)xq + (dim / 8) * index;
+                        minhash_jaccard_knn_ny(cur_query, (const char*)xb, hash_element_length, hash_element_size, nb,
+                                               topk, bitset, cur_distances, cur_labels);
+                    } else {
+                        size_t u8_dim = dim / 8;
+                        auto cur_query = (const char*)xq + u8_dim * index;
+                        minhash_lsh_hit_ny(cur_query, (const char*)xb, u8_dim, band, nb, topk, bitset, cur_distances,
+                                           cur_labels);
+                    }
+                    break;
+                }
                 case faiss::METRIC_Jaccard: {
                     auto cur_query = (const uint8_t*)xq + (dim / 8) * index;
                     faiss::float_maxheap_array_t res = {size_t(1), size_t(topk), cur_labels, cur_distances};
@@ -483,6 +534,10 @@ BruteForce::RangeSearch(const DataSetPtr base_dataset, const DataSetPtr query_da
     float range_filter = cfg.range_filter.value();
 
     auto pool = ThreadPool::GetGlobalSearchThreadPool();
+    // some check for minhash metric
+    if (metric_str == metric::MHJACCARD) {
+        return expected<DataSetPtr>::Err(Status::not_implemented, "minhash not support range search.");
+    }
 
     std::vector<std::vector<int64_t>> result_id_array(nq);
     std::vector<std::vector<float>> result_dist_array(nq);
@@ -758,6 +813,12 @@ BruteForce::AnnIterator(const DataSetPtr base_dataset, const DataSetPtr query_da
         return expected<std::vector<IndexNode::IteratorPtr>>::Err(result.error(), result.what());
     }
 
+    // some check for minhash metric
+    if (metric_str == metric::MHJACCARD) {
+        return expected<std::vector<IndexNode::IteratorPtr>>::Err(Status::not_implemented,
+                                                                  "minhash does not support iterator.");
+    }
+
 #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT)
     // LCOV_EXCL_START
     std::shared_ptr<tracer::trace::Span> span = nullptr;