move make offset map to load stage (#134)

Signed-off-by: xianliang <[email protected]>
zilliztech · Oct 9, 2023 · 771e8f3 · 771e8f3
1 parent 3051449
commit 771e8f3
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/src/index/ivf/ivf.cc b/src/index/ivf/ivf.cc
@@ -278,7 +278,6 @@ IvfIndexNode<T>::Train(const DataSet& dataset, const Config& cfg) {
             qzr = new (std::nothrow) typename QuantizerT<T>::type(dim, metric.value());
             index = std::make_unique<faiss::IndexIVFFlat>(qzr, dim, nlist, metric.value(), is_cosine);
             index->train(rows, (const float*)data);
-            index->make_direct_map(true);
         }
         if constexpr (std::is_same<faiss::IndexIVFFlatCC, T>::value) {
             const IvfFlatCcConfig& ivf_flat_cc_cfg = static_cast<const IvfFlatCcConfig&>(cfg);
@@ -287,6 +286,7 @@ IvfIndexNode<T>::Train(const DataSet& dataset, const Config& cfg) {
             index = std::make_unique<faiss::IndexIVFFlatCC>(qzr, dim, nlist, ivf_flat_cc_cfg.ssize.value(),
                                                             metric.value(), is_cosine);
             index->train(rows, (const float*)data);
+            // ivfflat_cc has no serialize stage, make map at build stage
             index->make_direct_map(true, faiss::DirectMap::ConcurrentArray);
         }
         if constexpr (std::is_same<faiss::IndexIVFPQ, T>::value) {
@@ -296,7 +296,6 @@ IvfIndexNode<T>::Train(const DataSet& dataset, const Config& cfg) {
             qzr = new (std::nothrow) typename QuantizerT<T>::type(dim, metric.value());
             index = std::make_unique<faiss::IndexIVFPQ>(qzr, dim, nlist, ivf_pq_cfg.m.value(), nbits, metric.value());
             index->train(rows, (const float*)data);
-            index->make_direct_map(true);
         }
         if constexpr (std::is_same<faiss::IndexScaNN, T>::value) {
             const ScannConfig& scann_cfg = static_cast<const ScannConfig&>(cfg);
@@ -320,15 +319,13 @@ IvfIndexNode<T>::Train(const DataSet& dataset, const Config& cfg) {
             index = std::make_unique<faiss::IndexIVFScalarQuantizer>(qzr, dim, nlist, faiss::QuantizerType::QT_8bit,
                                                                      metric.value());
             index->train(rows, (const float*)data);
-            index->make_direct_map(true);
         }
         if constexpr (std::is_same<faiss::IndexBinaryIVF, T>::value) {
             const IvfBinConfig& ivf_bin_cfg = static_cast<const IvfBinConfig&>(cfg);
             auto nlist = MatchNlist(rows, ivf_bin_cfg.nlist.value());
             qzr = new (std::nothrow) typename QuantizerT<T>::type(dim, metric.value());
             index = std::make_unique<faiss::IndexBinaryIVF>(qzr, dim, nlist, metric.value());
             index->train(rows, (const uint8_t*)data);
-            index->make_direct_map(true);
         }
         index->own_fields = true;
     } catch (std::exception& e) {
@@ -751,6 +748,12 @@ IvfIndexNode<T>::Deserialize(const BinarySet& binset, const Config& config) {
         } else {
             index_.reset(static_cast<T*>(faiss::read_index(&reader)));
         }
+        if constexpr (!std::is_same_v<T, faiss::IndexScaNN>) {
+            const BaseConfig& base_cfg = static_cast<const BaseConfig&>(config);
+            if (HasRawData(base_cfg.metric_type.value())) {
+                index_->make_direct_map(true);
+            }
+        }
     } catch (const std::exception& e) {
         LOG_KNOWHERE_WARNING_ << "faiss inner error: " << e.what();
         return Status::faiss_inner_error;
@@ -773,6 +776,12 @@ IvfIndexNode<T>::DeserializeFromFile(const std::string& filename, const Config&
         } else {
             index_.reset(static_cast<T*>(faiss::read_index(filename.data(), io_flags)));
         }
+        if constexpr (!std::is_same_v<T, faiss::IndexScaNN>) {
+            const BaseConfig& base_cfg = static_cast<const BaseConfig&>(config);
+            if (HasRawData(base_cfg.metric_type.value())) {
+                index_->make_direct_map(true);
+            }
+        }
     } catch (const std::exception& e) {
         LOG_KNOWHERE_WARNING_ << "faiss inner error: " << e.what();
         return Status::faiss_inner_error;