feat: add hnsw-rabitq

Author: egolearner

python/zvec/__init__.py

@@ -44,6 +44,7 @@
     FlatIndexParam,
     HnswIndexParam,
     HnswQueryParam,
+    HnswRabitqIndexParam,
     IndexOption,
     InvertIndexParam,
     IVFIndexParam,
@@ -90,6 +91,7 @@
     "VectorQuery",
     "InvertIndexParam",
     "HnswIndexParam",
+    "HnswRabitqIndexParam",
     "FlatIndexParam",
     "IVFIndexParam",
     "CollectionOption",

python/zvec/model/param/__init__.py

@@ -20,6 +20,7 @@
     FlatIndexParam,
     HnswIndexParam,
     HnswQueryParam,
+    HnswRabitqIndexParam,
     IndexOption,
     InvertIndexParam,
     IVFIndexParam,
@@ -34,6 +35,7 @@
     "FlatIndexParam",
     "HnswIndexParam",
     "HnswQueryParam",
+    "HnswRabitqIndexParam",
     "IVFIndexParam",
     "IVFQueryParam",
     "IndexOption",

src/binding/python/CMakeLists.txt

@@ -24,6 +24,7 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
             $<TARGET_FILE:core_knn_flat_static>
             $<TARGET_FILE:core_knn_flat_sparse_static>
             $<TARGET_FILE:core_knn_hnsw_static>
+            $<TARGET_FILE:core_knn_hnsw_rabitq_static>
             $<TARGET_FILE:core_knn_hnsw_sparse_static>
             $<TARGET_FILE:core_knn_ivf_static>
             $<TARGET_FILE:core_knn_cluster_static>
@@ -42,6 +43,7 @@ elseif (APPLE)
             -Wl,-force_load,$<TARGET_FILE:core_knn_flat_static>
             -Wl,-force_load,$<TARGET_FILE:core_knn_flat_sparse_static>
             -Wl,-force_load,$<TARGET_FILE:core_knn_hnsw_static>
+            -Wl,-force_load,$<TARGET_FILE:core_knn_hnsw_rabitq_static>
             -Wl,-force_load,$<TARGET_FILE:core_knn_hnsw_sparse_static>
             -Wl,-force_load,$<TARGET_FILE:core_knn_ivf_static>
             -Wl,-force_load,$<TARGET_FILE:core_knn_cluster_static>

src/binding/python/model/param/python_param.cc

@@ -55,6 +55,8 @@ static std::string quantize_type_to_string(const QuantizeType type) {
       return "UNDEFINED";
     case QuantizeType::INT8:
       return "INT8";
+    case QuantizeType::RABITQ:
+      return "RABITQ";
     case QuantizeType::INT4:
       return "INT4";
     case QuantizeType::FP16:
@@ -376,6 +378,43 @@ encapsulates its construction hyperparameters.
                 t[3].cast<QuantizeType>());
           }));
 
+  // binding hnsw rabitq index params
+  py::class_<HNSWRabitqIndexParams, VectorIndexParams,
+             std::shared_ptr<HNSWRabitqIndexParams>>
+      hnsw_rabitq_params(m, "HnswRabitqIndexParam", R"pbdoc(
+Parameters for configuring an HNSW (Hierarchical Navigable Small World) index with RabitQ.
+HNSW is a graph-based approximate nearest neighbor search index. This class
+encapsulates its construction hyperparameters.
+Attributes:
+    metric_type (MetricType): Distance metric used for similarity computation.
+        Default is ``MetricType.IP`` (inner product).
+    m (int): Number of bi-directional links created for every new element
+        during construction. Higher values improve accuracy but increase
+        memory usage and construction time. Default is 100.
+    ef_construction (int): Size of the dynamic candidate list for nearest
+        neighbors during index construction. Larger values yield better
+        graph quality at the cost of slower build time. Default is 500.
+    quantize_type (QuantizeType): Optional quantization type for vector
+        compression (e.g., FP16, INT8). Default is `QuantizeType.UNDEFINED` to
+        disable quantization.
+Examples:
+    >>> from zvec.typing import MetricType, QuantizeType
+    >>> params = HnswRabitqIndexParam(
+    ...     metric_type=MetricType.COSINE,
+    ...     m=16,
+    ...     ef_construction=200,
+    ...     quantize_type=QuantizeType.INT8
+    ... )
+    >>> print(params)
+    {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8'}
+)pbdoc");
+  hnsw_rabitq_params.def(
+      py::init<MetricType, int, int, QuantizeType>(),
+      py::arg("metric_type") = MetricType::IP,
+      py::arg("m") = core_interface::kDefaultHnswNeighborCnt,
+      py::arg("ef_construction") = core_interface::kDefaultHnswEfConstruction,
+      py::arg("quantize_type") = QuantizeType::UNDEFINED);
+
   // FlatIndexParams
   py::class_<FlatIndexParams, VectorIndexParams,
              std::shared_ptr<FlatIndexParams>>

src/binding/python/typing/python_type.cc

@@ -131,7 +131,8 @@ Enumeration of supported quantization types for vector compression.
       .value("UNDEFINED", QuantizeType::UNDEFINED)
       .value("FP16", QuantizeType::FP16)
       .value("INT8", QuantizeType::INT8)
-      .value("INT4", QuantizeType::INT4);
+      .value("INT4", QuantizeType::INT4)
+      .value("RABITQ", QuantizeType::RABITQ);
 }
 
 void ZVecPyTyping::bind_status(py::module_ &m) {

src/core/algorithm/CMakeLists.txt

@@ -6,4 +6,5 @@ cc_directory(flat)
 cc_directory(flat_sparse)
 cc_directory(ivf)
 cc_directory(hnsw)
-cc_directory(hnsw_sparse)
+cc_directory(hnsw_sparse)
+cc_directory(hnsw-rabitq)

src/core/algorithm/cluster/opt_kmeans_cluster.cc

@@ -1250,7 +1250,7 @@ int OptKmeansCluster::init(const IndexMeta &meta,
                            const ailego::Params &params) {
   auto type_ = meta.data_type();
 
-  if (meta.metric_name() == "InnerProduct") {
+  if (meta.metric_name() == "InnerProduct" || meta.metric_name() == "Cosine") {
     switch (type_) {
       case IndexMeta::DataType::DT_FP16: {
         algorithm_.reset(

src/core/algorithm/flat/flat_builder.cc

@@ -127,7 +127,6 @@ int FlatBuilder<BATCH_SIZE>::dump(const IndexDumper::Pointer &dumper) {
     return error_code;
   }
 
-  holder_ = nullptr;
   stats_.set_dumped_count(keys.size());
   stats_.set_dumped_costtime(stamp.milli_seconds());
   return 0;

src/core/algorithm/hnsw-rabitq/CMakeLists.txt

@@ -0,0 +1,11 @@
+include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
+include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
+
+cc_library(
+    NAME core_knn_hnsw_rabitq
+    STATIC SHARED STRICT ALWAYS_LINK
+    SRCS *.cc rabitq/*.cc
+    LIBS core_framework sparsehash rabitqlib
+    INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
+    VERSION "${PROXIMA_ZVEC_VERSION}"
+  )