diff --git a/benchmark/hdf5/benchmark_knowhere.h b/benchmark/hdf5/benchmark_knowhere.h index 6ead30098..0505d87e2 100644 --- a/benchmark/hdf5/benchmark_knowhere.h +++ b/benchmark/hdf5/benchmark_knowhere.h @@ -75,10 +75,12 @@ class Benchmark_knowhere : public Benchmark_hdf5 { } // IVFFLAT_NM should load raw data - knowhere::BinaryPtr bin = std::make_shared(); - bin->data = std::shared_ptr((uint8_t*)xb_); - bin->size = dim_ * nb_ * sizeof(float); - binary_set.Append("RAW_DATA", bin); + if (index_type_ == knowhere::IndexEnum::INDEX_FAISS_IVFFLAT && binary_set.GetByName("RAW_DATA") == nullptr) { + knowhere::BinaryPtr bin = std::make_shared(); + bin->data = std::shared_ptr((uint8_t*)xb_); + bin->size = dim_ * nb_ * sizeof(float); + binary_set.Append("RAW_DATA", bin); + } index.Deserialize(binary_set, conf); } diff --git a/include/knowhere/index_node.h b/include/knowhere/index_node.h index 2384abd09..5979071f9 100644 --- a/include/knowhere/index_node.h +++ b/include/knowhere/index_node.h @@ -18,11 +18,24 @@ #include "knowhere/dataset.h" #include "knowhere/expected.h" #include "knowhere/object.h" +#include "knowhere/version.h" namespace knowhere { class IndexNode : public Object { public: + IndexNode(const int32_t ver) : versoin_(ver) { + } + + IndexNode() : versoin_(Version::GetDefaultVersion()) { + } + + IndexNode(const IndexNode& other) : versoin_(other.versoin_) { + } + + IndexNode(const IndexNode&& other) : versoin_(other.versoin_) { + } + virtual Status Build(const DataSet& dataset, const Config& cfg) { RETURN_IF_ERROR(Train(dataset, cfg)); @@ -92,6 +105,9 @@ class IndexNode : public Object { virtual ~IndexNode() { } + + protected: + Version versoin_; }; } // namespace knowhere diff --git a/include/knowhere/utils.h b/include/knowhere/utils.h index b6b14e127..efc80d5cc 100644 --- a/include/knowhere/utils.h +++ b/include/knowhere/utils.h @@ -75,8 +75,7 @@ round_down(const T value, const T align) { } extern void -ConvertIVFFlatIfNeeded(const BinarySet& binset, const MetricType metric_type, const uint8_t* raw_data, - const size_t raw_size); +ConvertIVFFlat(const BinarySet& binset, const MetricType metric_type, const uint8_t* raw_data, const size_t raw_size); bool UseDiskLoad(const std::string& index_type, const int32_t& /*version*/); diff --git a/include/knowhere/version.h b/include/knowhere/version.h index 4c376997b..c4a82322c 100644 --- a/include/knowhere/version.h +++ b/include/knowhere/version.h @@ -21,12 +21,12 @@ namespace knowhere { namespace { static constexpr int32_t default_version = 0; static constexpr int32_t minimal_version = 0; -static constexpr int32_t current_version = 0; +static constexpr int32_t current_version = 1; } // namespace class Version { public: - explicit Version(const IndexVersion& version) : version_(version) { + explicit Version(const IndexVersion version) : version_(version) { } // used when version is not set @@ -59,7 +59,7 @@ class Version { // the version number IndexVersion - VersionNumber() { + VersionNumber() const { return version_; } diff --git a/src/common/utils.cc b/src/common/utils.cc index 51f4af687..c700d63ab 100644 --- a/src/common/utils.cc +++ b/src/common/utils.cc @@ -73,8 +73,7 @@ CopyAndNormalizeVecs(const float* x, size_t rows, int32_t dim) { } void -ConvertIVFFlatIfNeeded(const BinarySet& binset, const MetricType metric_type, const uint8_t* raw_data, - const size_t raw_size) { +ConvertIVFFlat(const BinarySet& binset, const MetricType metric_type, const uint8_t* raw_data, const size_t raw_size) { std::vector names = {"IVF", // compatible with knowhere-1.x knowhere::IndexEnum::INDEX_FAISS_IVFFLAT}; auto binary = binset.GetByNames(names); @@ -85,38 +84,23 @@ ConvertIVFFlatIfNeeded(const BinarySet& binset, const MetricType metric_type, co MemoryIOReader reader(binary->data.get(), binary->size); try { - uint32_t h; - reader.read(&h, sizeof(h), 1); - // only read IVF_FLAT index header - std::unique_ptr ivfl = std::make_unique(faiss::IndexIVFFlat()); - faiss::read_ivf_header(ivfl.get(), &reader); - ivfl->code_size = ivfl->d * sizeof(float); + std::unique_ptr ivfl; + ivfl.reset(static_cast(faiss::read_index_nm(&reader))); // is_cosine is not defined in IVF_FLAT_NM, so mark it from config ivfl->is_cosine = IsMetricType(metric_type, knowhere::metric::COSINE); - auto remains = binary->size - reader.tellg() - sizeof(uint32_t) - sizeof(ivfl->invlists->nlist) - - sizeof(ivfl->invlists->code_size); - auto invlist_size = sizeof(uint32_t) + sizeof(size_t) + ivfl->nlist * sizeof(size_t); - auto ids_size = ivfl->ntotal * sizeof(faiss::Index::idx_t); - // auto codes_size = ivfl->d * ivfl->ntotal * sizeof(float); - - // IVF_FLAT_NM format, need convert to new format - if (remains == invlist_size + ids_size) { - faiss::read_InvertedLists_nm(ivfl.get(), &reader); - ivfl->restore_codes(raw_data, raw_size); - - // over-write IVF_FLAT_NM binary with native IVF_FLAT binary - MemoryIOWriter writer; - faiss::write_index(ivfl.get(), &writer); - std::shared_ptr data(writer.data()); - binary->data = data; - binary->size = writer.tellg(); - - LOG_KNOWHERE_INFO_ << "Convert IVF_FLAT_NM to native IVF_FLAT, rows " << ivfl->ntotal << ", dim " - << ivfl->d; - } + ivfl->restore_codes(raw_data, raw_size); + + // over-write IVF_FLAT_NM binary with native IVF_FLAT binary + MemoryIOWriter writer; + faiss::write_index(ivfl.get(), &writer); + std::shared_ptr data(writer.data()); + binary->data = data; + binary->size = writer.tellg(); + + LOG_KNOWHERE_INFO_ << "Convert IVF_FLAT_NM to native IVF_FLAT, rows " << ivfl->ntotal << ", dim " << ivfl->d; } catch (...) { // not IVF_FLAT_NM format, do nothing return; diff --git a/src/index/flat/flat.cc b/src/index/flat/flat.cc index 1551b1baf..b6dd41b07 100644 --- a/src/index/flat/flat.cc +++ b/src/index/flat/flat.cc @@ -26,7 +26,7 @@ namespace knowhere { template class FlatIndexNode : public IndexNode { public: - FlatIndexNode(const int32_t& version, const Object& object) : index_(nullptr) { + FlatIndexNode(const int32_t version, const Object& object) : index_(nullptr) { static_assert(std::is_same::value || std::is_same::value, "not support"); search_pool_ = ThreadPool::GetGlobalSearchThreadPool(); @@ -237,7 +237,11 @@ class FlatIndexNode : public IndexNode { bool HasRawData(const std::string& metric_type) const override { if constexpr (std::is_same::value) { - return true; + if (versoin_ <= Version::GetMinimalVersion()) { + return !IsMetricType(metric_type, metric::COSINE); + } else { + return true; + } } if constexpr (std::is_same::value) { return true; diff --git a/src/index/ivf/ivf.cc b/src/index/ivf/ivf.cc index aceb669bd..f314bf4e8 100644 --- a/src/index/ivf/ivf.cc +++ b/src/index/ivf/ivf.cc @@ -45,7 +45,7 @@ struct QuantizerT { template class IvfIndexNode : public IndexNode { public: - IvfIndexNode(const int32_t& /*version*/, const Object& object) : index_(nullptr) { + IvfIndexNode(const int32_t version, const Object& object) : IndexNode(version), index_(nullptr) { static_assert(std::is_same::value || std::is_same::value || std::is_same::value || std::is_same::value || @@ -675,6 +675,50 @@ IvfIndexNode::Serialize(BinarySet& binset) const { } } +template <> +Status +IvfIndexNode::Serialize(BinarySet& binset) const { + try { + MemoryIOWriter writer; + LOG_KNOWHERE_INFO_ << "request version " << versoin_.VersionNumber(); + if (versoin_ <= Version::GetMinimalVersion()) { + faiss::write_index_nm(index_.get(), &writer); + LOG_KNOWHERE_INFO_ << "write IVF_FLAT_NM, file size " << writer.tellg(); + } else { + faiss::write_index(index_.get(), &writer); + LOG_KNOWHERE_INFO_ << "write IVF_FLAT, file size " << writer.tellg(); + } + std::shared_ptr index_data_ptr(writer.data()); + binset.Append(Type(), index_data_ptr, writer.tellg()); + + // append raw data for backward compatible + if (versoin_ <= Version::GetMinimalVersion()) { + size_t dim = index_->d; + size_t rows = index_->ntotal; + size_t raw_data_size = dim * rows * sizeof(float); + uint8_t* raw_data = new uint8_t[raw_data_size]; + std::shared_ptr raw_data_ptr(raw_data); + for (size_t i = 0; i < index_->nlist; i++) { + size_t list_size = index_->invlists->list_size(i); + const faiss::idx_t* ids = index_->invlists->get_ids(i); + const uint8_t* codes = index_->invlists->get_codes(i); + for (size_t j = 0; j < list_size; j++) { + faiss::idx_t id = ids[j]; + const uint8_t* src = codes + j * dim * sizeof(float); + uint8_t* dst = raw_data + id * dim * sizeof(float); + memcpy(dst, src, dim * sizeof(float)); + } + } + binset.Append("RAW_DATA", raw_data_ptr, raw_data_size); + LOG_KNOWHERE_INFO_ << "append raw data for IVF_FLAT_NM, size " << raw_data_size; + } + return Status::success; + } catch (const std::exception& e) { + LOG_KNOWHERE_WARNING_ << "faiss inner error: " << e.what(); + return Status::faiss_inner_error; + } +} + template Status IvfIndexNode::Deserialize(const BinarySet& binset, const Config& config) { @@ -690,10 +734,10 @@ IvfIndexNode::Deserialize(const BinarySet& binset, const Config& config) { MemoryIOReader reader(binary->data.get(), binary->size); try { if constexpr (std::is_same::value) { - auto raw_binary = binset.GetByName("RAW_DATA"); - if (raw_binary != nullptr) { + if (versoin_ <= Version::GetMinimalVersion()) { + auto raw_binary = binset.GetByName("RAW_DATA"); const BaseConfig& base_cfg = static_cast(config); - ConvertIVFFlatIfNeeded(binset, base_cfg.metric_type.value(), raw_binary->data.get(), raw_binary->size); + ConvertIVFFlat(binset, base_cfg.metric_type.value(), raw_binary->data.get(), raw_binary->size); // after conversion, binary size and data will be updated reader.data_ = binary->data.get(); reader.total_ = binary->size; diff --git a/thirdparty/faiss/faiss/impl/index_read.cpp b/thirdparty/faiss/faiss/impl/index_read.cpp index 6ff36885c..47324c591 100644 --- a/thirdparty/faiss/faiss/impl/index_read.cpp +++ b/thirdparty/faiss/faiss/impl/index_read.cpp @@ -551,10 +551,10 @@ static void read_direct_map(DirectMap* dm, IOReader* f) { } } -void read_ivf_header( +static void read_ivf_header( IndexIVF* ivf, IOReader* f, - std::vector>* ids) { + std::vector>* ids = nullptr) { read_index_header(ivf, f); READ1(ivf->nlist); READ1(ivf->nprobe); diff --git a/thirdparty/faiss/faiss/impl/index_write.cpp b/thirdparty/faiss/faiss/impl/index_write.cpp index efe674dd2..16ac8e370 100644 --- a/thirdparty/faiss/faiss/impl/index_write.cpp +++ b/thirdparty/faiss/faiss/impl/index_write.cpp @@ -376,6 +376,60 @@ void write_InvertedLists(const InvertedLists* ils, IOWriter* f) { } } +// write inverted lists for offset-only index +void write_InvertedLists_nm(const InvertedLists *ils, IOWriter *f) { + if (ils == nullptr) { + uint32_t h = fourcc("il00"); + WRITE1(h); + } else if (const auto & ails = + dynamic_cast(ils)) { + uint32_t h = fourcc("ilar"); + WRITE1(h); + WRITE1(ails->nlist); + WRITE1(ails->code_size); + // here we store either as a full or a sparse data buffer + size_t n_non0 = 0; + for (size_t i = 0; i < ails->nlist; i++) { + if (ails->ids[i].size() > 0) + n_non0++; + } + if (n_non0 > ails->nlist / 2) { + uint32_t list_type = fourcc("full"); + WRITE1(list_type); + std::vector sizes; + for (size_t i = 0; i < ails->nlist; i++) { + sizes.push_back (ails->ids[i].size()); + } + WRITEVECTOR(sizes); + } else { + int list_type = fourcc("sprs"); // sparse + WRITE1(list_type); + std::vector sizes; + for (size_t i = 0; i < ails->nlist; i++) { + size_t n = ails->ids[i].size(); + if (n > 0) { + sizes.push_back (i); + sizes.push_back (n); + } + } + WRITEVECTOR(sizes); + } + // make a single contiguous data buffer (useful for mmapping) + for (size_t i = 0; i < ails->nlist; i++) { + size_t n = ails->ids[i].size(); + if (n > 0) { + // WRITEANDCHECK (ails->codes[i].data(), n * ails->code_size); + WRITEANDCHECK(ails->ids[i].data(), n); + } + } + } else { + fprintf(stderr, "WARN! write_InvertedLists: unsupported invlist type, " + "saving null invlist\n"); + uint32_t h = fourcc("il00"); + WRITE1(h); + } +} + void write_ProductQuantizer(const ProductQuantizer* pq, const char* fname) { FileIOWriter writer(fname); write_ProductQuantizer(pq, &writer); @@ -729,6 +783,19 @@ void write_index(const Index* idx, const char* fname) { write_index(idx, &writer); } +// write index for offset-only index +void write_index_nm(const Index *idx, IOWriter *f) { + if(const IndexIVFFlat * ivfl = + dynamic_cast (idx)) { + uint32_t h = fourcc("IwFl"); + WRITE1(h); + write_ivf_header(ivfl, f); + write_InvertedLists_nm(ivfl->invlists, f); + } else { + FAISS_THROW_MSG("don't know how to serialize this type of index"); + } +} + void write_VectorTransform(const VectorTransform* vt, const char* fname) { FileIOWriter writer(fname); write_VectorTransform(vt, &writer); diff --git a/thirdparty/faiss/faiss/index_io.h b/thirdparty/faiss/faiss/index_io.h index 5a7009cdb..712c90470 100644 --- a/thirdparty/faiss/faiss/index_io.h +++ b/thirdparty/faiss/faiss/index_io.h @@ -79,9 +79,8 @@ void write_InvertedLists(const InvertedLists* ils, IOWriter* f); InvertedLists* read_InvertedLists(IOReader* reader, int io_flags = 0); // for backward compatibility -void read_ivf_header(IndexIVF* ivf, IOReader* f, - std::vector>* ids = nullptr); -void read_InvertedLists_nm(IndexIVF *ivf, IOReader *f, int io_flags = 0); +Index *read_index_nm(IOReader *f, int io_flags = 0); +void write_index_nm(const Index* idx, IOWriter* writer); } // namespace faiss #endif