Add nbits for SQ

Author: CLiqing

src/index/ivf/ivf.cc

@@ -614,6 +614,7 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
     if constexpr (std::is_same<IndexIVFSQWrapper, IndexType>::value) {
         const IvfSqConfig& ivf_sq_cfg = static_cast<const IvfSqConfig&>(*cfg);
         auto nlist = MatchNlist(rows, ivf_sq_cfg.nlist.value());
+        auto nbits = MatchNbits(rows, ivf_sq_cfg.nbits.value());
 
         const bool use_elkan = ivf_sq_cfg.use_elkan.value_or(true);
 
@@ -623,7 +624,7 @@ IvfIndexNode<DataType, IndexType>::TrainInternal(const DataSetPtr dataset, std::
 
         DataFormatEnum data_format = DataType2EnumHelper<DataType>::value;
 
-        auto result = IndexIvfFactory::create_for_sq(qzr.get(), dim, nlist, ivf_sq_cfg, data_format, metric.value());
+        auto result = IndexIvfFactory::create_for_sq(qzr.get(), dim, nlist, nbits, ivf_sq_cfg, data_format, metric.value());
         if (!result.has_value()) {
             return result.error();
         }

src/index/ivf/ivf_config.h

@@ -230,7 +230,35 @@ class ScannConfig : public IvfFlatConfig {
     }
 };
 
-class IvfSqConfig : public IvfConfig {};
+class IvfSqConfig : public IvfConfig {
+ public:
+    CFG_INT nbits;
+    KNOHWERE_DECLARE_CONFIG(IvfSqConfig) {
+        KNOWHERE_CONFIG_DECLARE_FIELD(nbits).description("nbits").set_default(8).for_train().set_range(4, 8);
+    }
+
+    Status
+    CheckAndAdjust(PARAM_TYPE param_type, std::string* err_msg) override {
+        // check the base class
+        const auto base_status = IvfConfig::CheckAndAdjust(param_type, err_msg);
+        if (base_status != Status::success) {
+            return base_status;
+        }
+
+        // check our parameters
+        if (param_type == PARAM_TYPE::TRAIN) {
+            // check nbits
+            if (nbits.has_value()) {
+                if (nbits.value() != 4 && nbits.value() != 6 && nbits.value() != 8) {
+                    std::string msg = "invalid nbits : " + std::to_string(nbits.value()) +
+                                      ", optional values are [4, 6, 8]";
+                    return HandleError(err_msg, msg, Status::invalid_args);
+                }
+            }
+        }
+        return Status::success;
+    }
+};
 
 class IvfBinConfig : public IvfConfig {
     Status

src/index/ivf/ivf_wrapper.cc

@@ -239,15 +239,23 @@ IndexIvfFactory::create_for_pq(faiss::IndexFlat* qzr_raw_ptr,
 
 expected<std::unique_ptr<IndexIVFSQWrapper>>
 IndexIvfFactory::create_for_sq(faiss::IndexFlat* qzr_raw_ptr,
-                               const faiss::idx_t d, const size_t nlist, const IvfSqConfig& ivf_sq_cfg,
+                               const faiss::idx_t d, const size_t nlist, const size_t nbits, const IvfSqConfig& ivf_sq_cfg,
                                const DataFormatEnum raw_data_format, const faiss::MetricType metric) {
     // the index factory string is either `IVFx,SQ,Refine(y)` or `IVFx,SQ`,
     //   depends on the refine parameters
 
     // create IndexIVFSQ
     // Index does not own qzr
+    faiss::ScalarQuantizer::QuantizerType quantizer_type;
+    if (nbits == 4) {
+        quantizer_type = faiss::ScalarQuantizer::QuantizerType::QT_4bit;
+    } else if (nbits == 6) {
+        quantizer_type = faiss::ScalarQuantizer::QuantizerType::QT_6bit;
+    } else {
+        quantizer_type = faiss::ScalarQuantizer::QuantizerType::QT_8bit;
+    }
     auto index = std::make_unique<faiss::IndexIVFScalarQuantizer>(
-        qzr_raw_ptr, d, nlist, faiss::ScalarQuantizer::QuantizerType::QT_8bit, metric);
+        qzr_raw_ptr, d, nlist, quantizer_type, metric);
 
     // create a refiner index, if needed
     std::unique_ptr<faiss::Index> idx_final;

src/index/ivf/ivf_wrapper.h

@@ -96,7 +96,7 @@ using IndexIVFPQWrapper = IndexIVFWrapper<faiss::IndexIVFPQ>;
 using IndexIVFSQWrapper = IndexIVFWrapper<faiss::IndexIVFScalarQuantizer>;
 
 class IndexIvfFactory {
-public:
+ public:
     static expected<std::unique_ptr<IndexIVFPQWrapper>>
     create_for_pq(faiss::IndexFlat* qzr_raw_ptr,
                   const faiss::idx_t d, const size_t nlist, const size_t nbits, const IvfPqConfig& ivf_pq_cfg,
@@ -105,7 +105,7 @@ class IndexIvfFactory {
 
     static expected<std::unique_ptr<IndexIVFSQWrapper>>
     create_for_sq(faiss::IndexFlat* qzr_raw_ptr,
-                  const faiss::idx_t d, const size_t nlist, const IvfSqConfig& ivf_sq_cfg,
+                  const faiss::idx_t d, const size_t nlist, const size_t nbits, const IvfSqConfig& ivf_sq_cfg,
                   // this is the data format of the raw data (if the refine is used)
                   const DataFormatEnum raw_data_format, const faiss::MetricType metric = faiss::METRIC_L2);
 };