From 9e43c45de436e093b37bddcf1ceb9caa7e323e01 Mon Sep 17 00:00:00 2001 From: Xu Peng Date: Tue, 7 May 2019 14:38:39 +0800 Subject: [PATCH] feat(db): support trait for engine Former-commit-id: cfb90e35386177032fa0c4ac7b8c81f29a536571 --- cpp/src/db/DBImpl.cpp | 3 +- cpp/src/db/FaissExecutionEngine.cpp | 44 ++++++++++++++++++++--------- cpp/src/db/FaissExecutionEngine.h | 10 ++++--- cpp/src/db/Traits.cpp | 19 +++++++++++++ cpp/src/db/Traits.h | 26 +++++++++++++++++ 5 files changed, 84 insertions(+), 18 deletions(-) create mode 100644 cpp/src/db/Traits.cpp create mode 100644 cpp/src/db/Traits.h diff --git a/cpp/src/db/DBImpl.cpp b/cpp/src/db/DBImpl.cpp index 70ad0b53dd..10d1d58463 100644 --- a/cpp/src/db/DBImpl.cpp +++ b/cpp/src/db/DBImpl.cpp @@ -17,6 +17,7 @@ #include "DBMetaImpl.h" #include "Env.h" #include "FaissExecutionEngine.h" +#include "Traits.h" namespace zilliz { namespace vecwise { @@ -415,7 +416,7 @@ DB::~DB() {} void DB::Open(const Options& options, DB** dbptr) { *dbptr = nullptr; - *dbptr = new DBImpl(options); + *dbptr = new DBImpl>(options); return; } diff --git a/cpp/src/db/FaissExecutionEngine.cpp b/cpp/src/db/FaissExecutionEngine.cpp index 06b7127217..76847705d3 100644 --- a/cpp/src/db/FaissExecutionEngine.cpp +++ b/cpp/src/db/FaissExecutionEngine.cpp @@ -3,6 +3,9 @@ * Unauthorized copying of this file, via any medium is strictly prohibited. * Proprietary and confidential. ******************************************************************************/ +#ifndef FAISSEXECUTIONENGINE_CPP__ +#define FAISSEXECUTIONENGINE_CPP__ + #include #include #include @@ -22,39 +25,47 @@ const std::string RawIndexType = "IDMap,Flat"; const std::string BuildIndexType = "IVF"; // IDMap / IVF -FaissExecutionEngine::FaissExecutionEngine(uint16_t dimension, const std::string& location) +template +FaissExecutionEngine::FaissExecutionEngine(uint16_t dimension, const std::string& location) : pIndex_(faiss::index_factory(dimension, RawIndexType.c_str())), location_(location) { } -FaissExecutionEngine::FaissExecutionEngine(std::shared_ptr index, const std::string& location) +template +FaissExecutionEngine::FaissExecutionEngine(std::shared_ptr index, const std::string& location) : pIndex_(index), location_(location) { } -Status FaissExecutionEngine::AddWithIds(long n, const float *xdata, const long *xids) { +template +Status FaissExecutionEngine::AddWithIds(long n, const float *xdata, const long *xids) { pIndex_->add_with_ids(n, xdata, xids); return Status::OK(); } -size_t FaissExecutionEngine::Count() const { +template +size_t FaissExecutionEngine::Count() const { return (size_t)(pIndex_->ntotal); } -size_t FaissExecutionEngine::Size() const { +template +size_t FaissExecutionEngine::Size() const { return (size_t)(Count() * pIndex_->d); } -size_t FaissExecutionEngine::PhysicalSize() const { +template +size_t FaissExecutionEngine::PhysicalSize() const { return (size_t)(Size()*sizeof(float)); } -Status FaissExecutionEngine::Serialize() { +template +Status FaissExecutionEngine::Serialize() { write_index(pIndex_.get(), location_.c_str()); return Status::OK(); } -Status FaissExecutionEngine::Load() { +template +Status FaissExecutionEngine::Load() { auto index = zilliz::vecwise::cache::CpuCacheMgr::GetInstance()->GetIndex(location_); if (!index) { index = read_index(location_); @@ -66,7 +77,8 @@ Status FaissExecutionEngine::Load() { return Status::OK(); } -Status FaissExecutionEngine::Merge(const std::string& location) { +template +Status FaissExecutionEngine::Merge(const std::string& location) { if (location == location_) { return Status::Error("Cannot Merge Self"); } @@ -80,7 +92,9 @@ Status FaissExecutionEngine::Merge(const std::string& location) { return Status::OK(); } -FaissExecutionEngine::Ptr FaissExecutionEngine::BuildIndex(const std::string& location) { +template +typename FaissExecutionEngine::Ptr +FaissExecutionEngine::BuildIndex(const std::string& location) { auto opd = std::make_shared(); opd->d = pIndex_->d; opd->index_type = BuildIndexType; @@ -92,12 +106,13 @@ FaissExecutionEngine::Ptr FaissExecutionEngine::BuildIndex(const std::string& lo dynamic_cast(from_index->index)->xb.data(), from_index->id_map.data()); - Ptr new_ee(new FaissExecutionEngine(index->data(), location)); + Ptr new_ee(new FaissExecutionEngine(index->data(), location)); new_ee->Serialize(); return new_ee; } -Status FaissExecutionEngine::Search(long n, +template +Status FaissExecutionEngine::Search(long n, const float *data, long k, float *distances, @@ -107,7 +122,8 @@ Status FaissExecutionEngine::Search(long n, return Status::OK(); } -Status FaissExecutionEngine::Cache() { +template +Status FaissExecutionEngine::Cache() { zilliz::vecwise::cache::CpuCacheMgr::GetInstance( )->InsertItem(location_, std::make_shared(pIndex_)); @@ -118,3 +134,5 @@ Status FaissExecutionEngine::Cache() { } // namespace engine } // namespace vecwise } // namespace zilliz + +#endif diff --git a/cpp/src/db/FaissExecutionEngine.h b/cpp/src/db/FaissExecutionEngine.h index f34021b418..739e09efca 100644 --- a/cpp/src/db/FaissExecutionEngine.h +++ b/cpp/src/db/FaissExecutionEngine.h @@ -18,11 +18,11 @@ namespace zilliz { namespace vecwise { namespace engine { -class FaissExecutionEngine; -class FaissExecutionEngine : public ExecutionEngine { +template +class FaissExecutionEngine : public ExecutionEngine> { public: - typedef std::shared_ptr Ptr; + typedef std::shared_ptr> Ptr; FaissExecutionEngine(uint16_t dimension, const std::string& location); FaissExecutionEngine(std::shared_ptr index, const std::string& location); @@ -50,7 +50,7 @@ public: float *distances, long *labels) const; - FaissExecutionEngine::Ptr BuildIndex(const std::string&); + Ptr BuildIndex(const std::string&); Status Cache(); protected: @@ -62,3 +62,5 @@ protected: } // namespace engine } // namespace vecwise } // namespace zilliz + +#include "FaissExecutionEngine.cpp" diff --git a/cpp/src/db/Traits.cpp b/cpp/src/db/Traits.cpp new file mode 100644 index 0000000000..0d900b9212 --- /dev/null +++ b/cpp/src/db/Traits.cpp @@ -0,0 +1,19 @@ +/******************************************************************************* + * Copyright 上海赜睿信息科技有限公司(Zilliz) - All Rights Reserved + * Unauthorized copying of this file, via any medium is strictly prohibited. + * Proprietary and confidential. + ******************************************************************************/ + +namespace zilliz { +namespace vecwise { +namespace engine { + +const char* IVFIndexTrait::BuildIndexType = "IVF"; +const char* IVFIndexTrait::RawIndexType = "IDMap,Flat"; + +const char* IDMapIndexTrait::BuildIndexType = "IDMap"; +const char* IDMapIndexTrait::RawIndexType = "IDMap,Flat"; + +} // namespace engine +} // namespace vecwise +} // namespace zilliz diff --git a/cpp/src/db/Traits.h b/cpp/src/db/Traits.h new file mode 100644 index 0000000000..1740bc87ed --- /dev/null +++ b/cpp/src/db/Traits.h @@ -0,0 +1,26 @@ +/******************************************************************************* + * Copyright 上海赜睿信息科技有限公司(Zilliz) - All Rights Reserved + * Unauthorized copying of this file, via any medium is strictly prohibited. + * Proprietary and confidential. + ******************************************************************************/ +#pragma once + + +namespace zilliz { +namespace vecwise { +namespace engine { + +struct IVFIndexTrait { + static const char* BuildIndexType; + static const char* RawIndexType; +}; + +struct IDMapIndexTrait { + static const char* BuildIndexType; + static const char* RawIndexType; +}; + + +} // namespace engine +} // namespace vecwise +} // namespace zilliz