Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
BaiXuePrincess
milvus
提交
7633437e
milvus
项目概览
BaiXuePrincess
/
milvus
与 Fork 源项目一致
从无法访问的项目Fork
通知
7
Star
4
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
milvus
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
7633437e
编写于
9月 28, 2019
作者:
S
starlord
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
format wrapper code
Former-commit-id: 75fc82a41b150fbd3f1ddcc9d09c7c4b4e43c93b
上级
e9d3f80c
变更
13
隐藏空白更改
内联
并排
Showing
13 changed file
with
217 addition
and
269 deletion
+217
-269
cpp/build-support/lint_exclusions.txt
cpp/build-support/lint_exclusions.txt
+1
-2
cpp/src/wrapper/ConfAdapter.cpp
cpp/src/wrapper/ConfAdapter.cpp
+22
-22
cpp/src/wrapper/ConfAdapter.h
cpp/src/wrapper/ConfAdapter.h
+19
-17
cpp/src/wrapper/ConfAdapterMgr.cpp
cpp/src/wrapper/ConfAdapterMgr.cpp
+9
-11
cpp/src/wrapper/ConfAdapterMgr.h
cpp/src/wrapper/ConfAdapterMgr.h
+11
-16
cpp/src/wrapper/DataTransfer.cpp
cpp/src/wrapper/DataTransfer.cpp
+11
-12
cpp/src/wrapper/DataTransfer.h
cpp/src/wrapper/DataTransfer.h
+5
-6
cpp/src/wrapper/KnowhereResource.cpp
cpp/src/wrapper/KnowhereResource.cpp
+16
-17
cpp/src/wrapper/KnowhereResource.h
cpp/src/wrapper/KnowhereResource.h
+3
-4
cpp/src/wrapper/VecImpl.cpp
cpp/src/wrapper/VecImpl.cpp
+38
-49
cpp/src/wrapper/VecImpl.h
cpp/src/wrapper/VecImpl.h
+25
-38
cpp/src/wrapper/VecIndex.cpp
cpp/src/wrapper/VecIndex.cpp
+37
-42
cpp/src/wrapper/VecIndex.h
cpp/src/wrapper/VecIndex.h
+20
-33
未找到文件。
cpp/build-support/lint_exclusions.txt
浏览文件 @
7633437e
...
...
@@ -6,5 +6,4 @@
*easylogging++*
*SqliteMetaImpl.cpp
*src/grpc*
*src/core*
*src/wrapper*
\ No newline at end of file
*src/core*
\ No newline at end of file
cpp/src/wrapper/ConfAdapter.cpp
浏览文件 @
7633437e
...
...
@@ -15,13 +15,14 @@
// specific language governing permissions and limitations
// under the License.
#include "wrapper/ConfAdapter.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "utils/Log.h"
#include <cmath>
#include "ConfAdapter.h"
#include "src/utils/Log.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include <memory>
// TODO: add conf checker
// TODO
(lxj)
: add conf checker
namespace
zilliz
{
namespace
milvus
{
...
...
@@ -42,7 +43,7 @@ ConfAdapter::MatchBase(knowhere::Config conf) {
}
knowhere
::
Config
ConfAdapter
::
Match
(
const
TempMetaConf
&
metaconf
)
{
ConfAdapter
::
Match
(
const
TempMetaConf
&
metaconf
)
{
auto
conf
=
std
::
make_shared
<
knowhere
::
Cfg
>
();
conf
->
d
=
metaconf
.
dim
;
conf
->
metric_type
=
metaconf
.
metric_type
;
...
...
@@ -52,14 +53,14 @@ ConfAdapter::Match(const TempMetaConf &metaconf) {
}
knowhere
::
Config
ConfAdapter
::
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
)
{
ConfAdapter
::
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
)
{
auto
conf
=
std
::
make_shared
<
knowhere
::
Cfg
>
();
conf
->
k
=
metaconf
.
k
;
return
conf
;
}
knowhere
::
Config
IVFConfAdapter
::
Match
(
const
TempMetaConf
&
metaconf
)
{
IVFConfAdapter
::
Match
(
const
TempMetaConf
&
metaconf
)
{
auto
conf
=
std
::
make_shared
<
knowhere
::
IVFCfg
>
();
conf
->
nlist
=
MatchNlist
(
metaconf
.
size
,
metaconf
.
nlist
);
conf
->
d
=
metaconf
.
dim
;
...
...
@@ -72,7 +73,7 @@ IVFConfAdapter::Match(const TempMetaConf &metaconf) {
static
constexpr
float
TYPICAL_COUNT
=
1000000.0
;
int64_t
IVFConfAdapter
::
MatchNlist
(
const
int64_t
&
size
,
const
int64_t
&
nlist
)
{
IVFConfAdapter
::
MatchNlist
(
const
int64_t
&
size
,
const
int64_t
&
nlist
)
{
if
(
size
<=
TYPICAL_COUNT
/
16384
+
1
)
{
// handle less row count, avoid nlist set to 0
return
1
;
...
...
@@ -84,7 +85,7 @@ IVFConfAdapter::MatchNlist(const int64_t &size, const int64_t &nlist) {
}
knowhere
::
Config
IVFConfAdapter
::
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
)
{
IVFConfAdapter
::
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
)
{
auto
conf
=
std
::
make_shared
<
knowhere
::
IVFCfg
>
();
conf
->
k
=
metaconf
.
k
;
conf
->
nprobe
=
metaconf
.
nprobe
;
...
...
@@ -95,17 +96,16 @@ IVFConfAdapter::MatchSearch(const TempMetaConf &metaconf, const IndexType &type)
case
IndexType
::
FAISS_IVFPQ_GPU
:
if
(
conf
->
nprobe
>
GPU_MAX_NRPOBE
)
{
WRAPPER_LOG_WARNING
<<
"When search with GPU, nprobe shoud be no more than "
<<
GPU_MAX_NRPOBE
<<
", but you passed "
<<
conf
->
nprobe
<<
"
. Search with "
<<
GPU_MAX_NRPOBE
<<
"
instead"
;
<<
", but you passed "
<<
conf
->
nprobe
<<
". Search with "
<<
GPU_MAX_NRPOBE
<<
" instead"
;
conf
->
nprobe
=
GPU_MAX_NRPOBE
;
}
}
return
conf
;
}
knowhere
::
Config
IVFSQConfAdapter
::
Match
(
const
TempMetaConf
&
metaconf
)
{
IVFSQConfAdapter
::
Match
(
const
TempMetaConf
&
metaconf
)
{
auto
conf
=
std
::
make_shared
<
knowhere
::
IVFSQCfg
>
();
conf
->
nlist
=
MatchNlist
(
metaconf
.
size
,
metaconf
.
nlist
);
conf
->
d
=
metaconf
.
dim
;
...
...
@@ -117,7 +117,7 @@ IVFSQConfAdapter::Match(const TempMetaConf &metaconf) {
}
knowhere
::
Config
IVFPQConfAdapter
::
Match
(
const
TempMetaConf
&
metaconf
)
{
IVFPQConfAdapter
::
Match
(
const
TempMetaConf
&
metaconf
)
{
auto
conf
=
std
::
make_shared
<
knowhere
::
IVFPQCfg
>
();
conf
->
nlist
=
MatchNlist
(
metaconf
.
size
,
metaconf
.
nlist
);
conf
->
d
=
metaconf
.
dim
;
...
...
@@ -130,7 +130,7 @@ IVFPQConfAdapter::Match(const TempMetaConf &metaconf) {
}
knowhere
::
Config
NSGConfAdapter
::
Match
(
const
TempMetaConf
&
metaconf
)
{
NSGConfAdapter
::
Match
(
const
TempMetaConf
&
metaconf
)
{
auto
conf
=
std
::
make_shared
<
knowhere
::
NSGCfg
>
();
conf
->
nlist
=
MatchNlist
(
metaconf
.
size
,
metaconf
.
nlist
);
conf
->
d
=
metaconf
.
dim
;
...
...
@@ -146,20 +146,20 @@ NSGConfAdapter::Match(const TempMetaConf &metaconf) {
conf
->
candidate_pool_size
=
200
+
100
*
scale_factor
;
MatchBase
(
conf
);
// WRAPPER_LOG_DEBUG << "nlist: " << conf->nlist
// << ", gpu_id: " << conf->gpu_id << ", d: " << conf->d
// << ", nprobe: " << conf->nprobe << ", knng: " << conf->knng;
// WRAPPER_LOG_DEBUG << "nlist: " << conf->nlist
// << ", gpu_id: " << conf->gpu_id << ", d: " << conf->d
// << ", nprobe: " << conf->nprobe << ", knng: " << conf->knng;
return
conf
;
}
knowhere
::
Config
NSGConfAdapter
::
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
)
{
NSGConfAdapter
::
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
)
{
auto
conf
=
std
::
make_shared
<
knowhere
::
NSGCfg
>
();
conf
->
k
=
metaconf
.
k
;
conf
->
search_length
=
metaconf
.
search_length
;
return
conf
;
}
}
}
}
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/ConfAdapter.h
浏览文件 @
7633437e
...
...
@@ -15,12 +15,12 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
#include "knowhere/common/Config.h"
#include "VecIndex.h"
#include "knowhere/common/Config.h"
#include <memory>
namespace
zilliz
{
namespace
milvus
{
...
...
@@ -42,16 +42,17 @@ struct TempMetaConf {
class
ConfAdapter
{
public:
virtual
knowhere
::
Config
Match
(
const
TempMetaConf
&
metaconf
);
Match
(
const
TempMetaConf
&
metaconf
);
virtual
knowhere
::
Config
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
);
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
);
// virtual void
// Dump(){}
// virtual void
// Dump(){}
protected:
static
void
MatchBase
(
knowhere
::
Config
conf
);
static
void
MatchBase
(
knowhere
::
Config
conf
);
};
using
ConfAdapterPtr
=
std
::
shared_ptr
<
ConfAdapter
>
;
...
...
@@ -59,36 +60,37 @@ using ConfAdapterPtr = std::shared_ptr<ConfAdapter>;
class
IVFConfAdapter
:
public
ConfAdapter
{
public:
knowhere
::
Config
Match
(
const
TempMetaConf
&
metaconf
)
override
;
Match
(
const
TempMetaConf
&
metaconf
)
override
;
knowhere
::
Config
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
)
override
;
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
)
override
;
protected:
static
int64_t
MatchNlist
(
const
int64_t
&
size
,
const
int64_t
&
nlist
);
static
int64_t
MatchNlist
(
const
int64_t
&
size
,
const
int64_t
&
nlist
);
};
class
IVFSQConfAdapter
:
public
IVFConfAdapter
{
public:
knowhere
::
Config
Match
(
const
TempMetaConf
&
metaconf
)
override
;
Match
(
const
TempMetaConf
&
metaconf
)
override
;
};
class
IVFPQConfAdapter
:
public
IVFConfAdapter
{
public:
knowhere
::
Config
Match
(
const
TempMetaConf
&
metaconf
)
override
;
Match
(
const
TempMetaConf
&
metaconf
)
override
;
};
class
NSGConfAdapter
:
public
IVFConfAdapter
{
public:
knowhere
::
Config
Match
(
const
TempMetaConf
&
metaconf
)
override
;
Match
(
const
TempMetaConf
&
metaconf
)
override
;
knowhere
::
Config
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
)
final
;
MatchSearch
(
const
TempMetaConf
&
metaconf
,
const
IndexType
&
type
)
final
;
};
}
}
}
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/ConfAdapterMgr.cpp
浏览文件 @
7633437e
...
...
@@ -15,18 +15,17 @@
// specific language governing permissions and limitations
// under the License.
#include "src/utils/Exception.h"
#include "ConfAdapterMgr.h"
#include "wrapper/ConfAdapterMgr.h"
#include "utils/Exception.h"
namespace
zilliz
{
namespace
milvus
{
namespace
engine
{
ConfAdapterPtr
AdapterMgr
::
GetAdapter
(
const
IndexType
&
indexType
)
{
if
(
!
init_
)
RegisterAdapter
();
AdapterMgr
::
GetAdapter
(
const
IndexType
&
indexType
)
{
if
(
!
init_
)
RegisterAdapter
();
auto
it
=
table_
.
find
(
indexType
);
if
(
it
!=
table_
.
end
())
{
...
...
@@ -36,8 +35,8 @@ AdapterMgr::GetAdapter(const IndexType &indexType) {
}
}
#define REGISTER_CONF_ADAPTER(T, KEY, NAME) static AdapterMgr::register_t<T> reg_##NAME##_(KEY)
#define REGISTER_CONF_ADAPTER(T, KEY, NAME) static AdapterMgr::register_t<T>reg_##NAME##_(KEY)
void
AdapterMgr
::
RegisterAdapter
()
{
init_
=
true
;
...
...
@@ -58,7 +57,6 @@ AdapterMgr::RegisterAdapter() {
REGISTER_CONF_ADAPTER
(
NSGConfAdapter
,
IndexType
::
NSG_MIX
,
nsg_mix
);
}
}
// engine
}
// milvus
}
// zilliz
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/ConfAdapterMgr.h
浏览文件 @
7633437e
...
...
@@ -15,12 +15,13 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
#include "VecIndex.h"
#include "ConfAdapter.h"
#include "VecIndex.h"
#include <map>
#include <memory>
namespace
zilliz
{
namespace
milvus
{
...
...
@@ -28,23 +29,21 @@ namespace engine {
class
AdapterMgr
{
public:
template
<
typename
T
>
template
<
typename
T
>
struct
register_t
{
explicit
register_t
(
const
IndexType
&
key
)
{
AdapterMgr
::
GetInstance
().
table_
.
emplace
(
key
,
[]
{
return
std
::
make_shared
<
T
>
();
});
explicit
register_t
(
const
IndexType
&
key
)
{
AdapterMgr
::
GetInstance
().
table_
.
emplace
(
key
,
[]
{
return
std
::
make_shared
<
T
>
();
});
}
};
static
AdapterMgr
&
static
AdapterMgr
&
GetInstance
()
{
static
AdapterMgr
instance
;
return
instance
;
}
ConfAdapterPtr
GetAdapter
(
const
IndexType
&
indexType
);
GetAdapter
(
const
IndexType
&
indexType
);
void
RegisterAdapter
();
...
...
@@ -54,10 +53,6 @@ class AdapterMgr {
std
::
map
<
IndexType
,
std
::
function
<
ConfAdapterPtr
()
>
>
table_
;
};
}
// engine
}
// milvus
}
// zilliz
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/DataTransfer.cpp
浏览文件 @
7633437e
...
...
@@ -15,39 +15,38 @@
// specific language governing permissions and limitations
// under the License.
#include "wrapper/DataTransfer.h"
#include <vector>
#include <memory>
#include <utility>
#include <vector>
namespace
zilliz
{
namespace
milvus
{
namespace
engine
{
knowhere
::
DatasetPtr
GenDatasetWithIds
(
const
int64_t
&
nb
,
const
int64_t
&
dim
,
const
float
*
xb
,
const
int64_t
*
ids
)
{
GenDatasetWithIds
(
const
int64_t
&
nb
,
const
int64_t
&
dim
,
const
float
*
xb
,
const
int64_t
*
ids
)
{
std
::
vector
<
int64_t
>
shape
{
nb
,
dim
};
auto
tensor
=
knowhere
::
ConstructFloatTensor
((
uint8_t
*
)
xb
,
nb
*
dim
*
sizeof
(
float
),
shape
);
auto
tensor
=
knowhere
::
ConstructFloatTensor
((
uint8_t
*
)
xb
,
nb
*
dim
*
sizeof
(
float
),
shape
);
std
::
vector
<
knowhere
::
TensorPtr
>
tensors
{
tensor
};
std
::
vector
<
knowhere
::
FieldPtr
>
tensor_fields
{
knowhere
::
ConstructFloatField
(
"data"
)};
auto
tensor_schema
=
std
::
make_shared
<
knowhere
::
Schema
>
(
tensor_fields
);
auto
id_array
=
knowhere
::
ConstructInt64Array
((
uint8_t
*
)
ids
,
nb
*
sizeof
(
int64_t
));
auto
id_array
=
knowhere
::
ConstructInt64Array
((
uint8_t
*
)
ids
,
nb
*
sizeof
(
int64_t
));
std
::
vector
<
knowhere
::
ArrayPtr
>
arrays
{
id_array
};
std
::
vector
<
knowhere
::
FieldPtr
>
array_fields
{
knowhere
::
ConstructInt64Field
(
"id"
)};
auto
array_schema
=
std
::
make_shared
<
knowhere
::
Schema
>
(
tensor_fields
);
auto
dataset
=
std
::
make_shared
<
knowhere
::
Dataset
>
(
std
::
move
(
arrays
),
array_schema
,
std
::
move
(
tensors
),
tensor_schema
);
auto
dataset
=
std
::
make_shared
<
knowhere
::
Dataset
>
(
std
::
move
(
arrays
),
array_schema
,
std
::
move
(
tensors
),
tensor_schema
);
return
dataset
;
}
knowhere
::
DatasetPtr
GenDataset
(
const
int64_t
&
nb
,
const
int64_t
&
dim
,
const
float
*
xb
)
{
GenDataset
(
const
int64_t
&
nb
,
const
int64_t
&
dim
,
const
float
*
xb
)
{
std
::
vector
<
int64_t
>
shape
{
nb
,
dim
};
auto
tensor
=
knowhere
::
ConstructFloatTensor
((
uint8_t
*
)
xb
,
nb
*
dim
*
sizeof
(
float
),
shape
);
auto
tensor
=
knowhere
::
ConstructFloatTensor
((
uint8_t
*
)
xb
,
nb
*
dim
*
sizeof
(
float
),
shape
);
std
::
vector
<
knowhere
::
TensorPtr
>
tensors
{
tensor
};
std
::
vector
<
knowhere
::
FieldPtr
>
tensor_fields
{
knowhere
::
ConstructFloatField
(
"data"
)};
auto
tensor_schema
=
std
::
make_shared
<
knowhere
::
Schema
>
(
tensor_fields
);
...
...
@@ -56,6 +55,6 @@ GenDataset(const int64_t &nb, const int64_t &dim, const float *xb) {
return
dataset
;
}
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/DataTransfer.h
浏览文件 @
7633437e
...
...
@@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
#include "knowhere/adapter/Structure.h"
...
...
@@ -25,11 +24,11 @@ namespace milvus {
namespace
engine
{
extern
zilliz
::
knowhere
::
DatasetPtr
GenDatasetWithIds
(
const
int64_t
&
nb
,
const
int64_t
&
dim
,
const
float
*
xb
,
const
int64_t
*
ids
);
GenDatasetWithIds
(
const
int64_t
&
nb
,
const
int64_t
&
dim
,
const
float
*
xb
,
const
int64_t
*
ids
);
extern
zilliz
::
knowhere
::
DatasetPtr
GenDataset
(
const
int64_t
&
nb
,
const
int64_t
&
dim
,
const
float
*
xb
);
GenDataset
(
const
int64_t
&
nb
,
const
int64_t
&
dim
,
const
float
*
xb
);
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/KnowhereResource.cpp
浏览文件 @
7633437e
...
...
@@ -15,16 +15,15 @@
// specific language governing permissions and limitations
// under the License.
#include "wrapper/KnowhereResource.h"
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
#include "server/Config.h"
#include <map>
#include <set>
#include <vector>
#include <string>
#include <utility>
#include <vector>
namespace
zilliz
{
namespace
milvus
{
...
...
@@ -43,22 +42,24 @@ KnowhereResource::Initialize() {
GpuResourcesArray
gpu_resources
;
Status
s
;
//get build index gpu resource
server
::
Config
&
config
=
server
::
Config
::
GetInstance
();
//
get build index gpu resource
server
::
Config
&
config
=
server
::
Config
::
GetInstance
();
int32_t
build_index_gpu
;
s
=
config
.
GetDBConfigBuildIndexGPU
(
build_index_gpu
);
if
(
!
s
.
ok
())
return
s
;
if
(
!
s
.
ok
())
return
s
;
gpu_resources
.
insert
(
std
::
make_pair
(
build_index_gpu
,
GpuResourceSetting
()));
//get search gpu resource
//
get search gpu resource
std
::
vector
<
std
::
string
>
pool
;
s
=
config
.
GetResourceConfigPool
(
pool
);
if
(
!
s
.
ok
())
return
s
;
if
(
!
s
.
ok
())
return
s
;
std
::
set
<
uint64_t
>
gpu_ids
;
for
(
auto
&
resource
:
pool
)
{
for
(
auto
&
resource
:
pool
)
{
if
(
resource
.
length
()
<
4
||
resource
.
substr
(
0
,
3
)
!=
"gpu"
)
{
// invalid
continue
;
...
...
@@ -67,12 +68,10 @@ KnowhereResource::Initialize() {
gpu_resources
.
insert
(
std
::
make_pair
(
gpu_id
,
GpuResourceSetting
()));
}
//init gpu resources
//
init gpu resources
for
(
auto
iter
=
gpu_resources
.
begin
();
iter
!=
gpu_resources
.
end
();
++
iter
)
{
knowhere
::
FaissGpuResourceMgr
::
GetInstance
().
InitDevice
(
iter
->
first
,
iter
->
second
.
pinned_memory
,
iter
->
second
.
temp_memory
,
iter
->
second
.
resource_num
);
knowhere
::
FaissGpuResourceMgr
::
GetInstance
().
InitDevice
(
iter
->
first
,
iter
->
second
.
pinned_memory
,
iter
->
second
.
temp_memory
,
iter
->
second
.
resource_num
);
}
return
Status
::
OK
();
...
...
@@ -80,10 +79,10 @@ KnowhereResource::Initialize() {
Status
KnowhereResource
::
Finalize
()
{
knowhere
::
FaissGpuResourceMgr
::
GetInstance
().
Free
();
// free gpu resource.
knowhere
::
FaissGpuResourceMgr
::
GetInstance
().
Free
();
// free gpu resource.
return
Status
::
OK
();
}
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/KnowhereResource.h
浏览文件 @
7633437e
...
...
@@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
#include "utils/Status.h"
...
...
@@ -33,6 +32,6 @@ class KnowhereResource {
Finalize
();
};
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/VecImpl.cpp
浏览文件 @
7633437e
...
...
@@ -15,14 +15,13 @@
// specific language governing permissions and limitations
// under the License.
#include "wrapper/VecImpl.h"
#include "utils/Log.h"
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "DataTransfer.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#include "
DataTransfer
.h"
#include "
utils/Log
.h"
/*
* no parameter check in this layer.
...
...
@@ -34,12 +33,8 @@ namespace milvus {
namespace
engine
{
Status
VecIndexImpl
::
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
{
VecIndexImpl
::
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
{
try
{
dim
=
cfg
->
d
;
auto
dataset
=
GenDatasetWithIds
(
nb
,
dim
,
xb
,
ids
);
...
...
@@ -49,10 +44,10 @@ VecIndexImpl::BuildAll(const int64_t &nb,
auto
model
=
index_
->
Train
(
dataset
,
cfg
);
index_
->
set_index_model
(
model
);
index_
->
Add
(
dataset
,
cfg
);
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_UNEXPECTED_ERROR
,
e
.
what
());
}
catch
(
std
::
exception
&
e
)
{
}
catch
(
std
::
exception
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_ERROR
,
e
.
what
());
}
...
...
@@ -60,15 +55,15 @@ VecIndexImpl::BuildAll(const int64_t &nb,
}
Status
VecIndexImpl
::
Add
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
)
{
VecIndexImpl
::
Add
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
)
{
try
{
auto
dataset
=
GenDatasetWithIds
(
nb
,
dim
,
xb
,
ids
);
index_
->
Add
(
dataset
,
cfg
);
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_UNEXPECTED_ERROR
,
e
.
what
());
}
catch
(
std
::
exception
&
e
)
{
}
catch
(
std
::
exception
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_ERROR
,
e
.
what
());
}
...
...
@@ -76,7 +71,7 @@ VecIndexImpl::Add(const int64_t &nb, const float *xb, const int64_t *ids, const
}
Status
VecIndexImpl
::
Search
(
const
int64_t
&
nq
,
const
float
*
xq
,
float
*
dist
,
int64_t
*
ids
,
const
Config
&
cfg
)
{
VecIndexImpl
::
Search
(
const
int64_t
&
nq
,
const
float
*
xq
,
float
*
dist
,
int64_t
*
ids
,
const
Config
&
cfg
)
{
try
{
auto
k
=
cfg
->
k
;
auto
dataset
=
GenDataset
(
nq
,
dim
,
xq
);
...
...
@@ -110,10 +105,10 @@ VecIndexImpl::Search(const int64_t &nq, const float *xq, float *dist, int64_t *i
// TODO(linxj): avoid copy here.
memcpy
(
ids
,
p_ids
,
sizeof
(
int64_t
)
*
nq
*
k
);
memcpy
(
dist
,
p_dist
,
sizeof
(
float
)
*
nq
*
k
);
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_UNEXPECTED_ERROR
,
e
.
what
());
}
catch
(
std
::
exception
&
e
)
{
}
catch
(
std
::
exception
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_ERROR
,
e
.
what
());
}
...
...
@@ -127,7 +122,7 @@ VecIndexImpl::Serialize() {
}
Status
VecIndexImpl
::
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
{
VecIndexImpl
::
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
{
index_
->
Load
(
index_binary
);
dim
=
Dimension
();
return
Status
::
OK
();
...
...
@@ -149,7 +144,7 @@ VecIndexImpl::GetType() {
}
VecIndexPtr
VecIndexImpl
::
CopyToGpu
(
const
int64_t
&
device_id
,
const
Config
&
cfg
)
{
VecIndexImpl
::
CopyToGpu
(
const
int64_t
&
device_id
,
const
Config
&
cfg
)
{
// TODO(linxj): exception handle
auto
gpu_index
=
zilliz
::
knowhere
::
cloner
::
CopyCpuToGpu
(
index_
,
device_id
,
cfg
);
auto
new_index
=
std
::
make_shared
<
VecIndexImpl
>
(
gpu_index
,
ConvertToGpuIndexType
(
type
));
...
...
@@ -158,7 +153,7 @@ VecIndexImpl::CopyToGpu(const int64_t &device_id, const Config &cfg) {
}
VecIndexPtr
VecIndexImpl
::
CopyToCpu
(
const
Config
&
cfg
)
{
VecIndexImpl
::
CopyToCpu
(
const
Config
&
cfg
)
{
// TODO(linxj): exception handle
auto
cpu_index
=
zilliz
::
knowhere
::
cloner
::
CopyGpuToCpu
(
index_
,
cfg
);
auto
new_index
=
std
::
make_shared
<
VecIndexImpl
>
(
cpu_index
,
ConvertToCpuIndexType
(
type
));
...
...
@@ -180,30 +175,32 @@ VecIndexImpl::GetDeviceId() {
return
device_idx
->
GetGpuDevice
();
}
// else
return
-
1
;
// -1 == cpu
return
-
1
;
// -1 == cpu
}
float
*
float
*
BFIndex
::
GetRawVectors
()
{
auto
raw_index
=
std
::
dynamic_pointer_cast
<
knowhere
::
IDMAP
>
(
index_
);
if
(
raw_index
)
{
return
raw_index
->
GetRawVectors
();
}
if
(
raw_index
)
{
return
raw_index
->
GetRawVectors
();
}
return
nullptr
;
}
int64_t
*
int64_t
*
BFIndex
::
GetRawIds
()
{
return
std
::
static_pointer_cast
<
knowhere
::
IDMAP
>
(
index_
)
->
GetRawIds
();
}
ErrorCode
BFIndex
::
Build
(
const
Config
&
cfg
)
{
BFIndex
::
Build
(
const
Config
&
cfg
)
{
try
{
dim
=
cfg
->
d
;
std
::
static_pointer_cast
<
knowhere
::
IDMAP
>
(
index_
)
->
Train
(
cfg
);
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
KNOWHERE_UNEXPECTED_ERROR
;
}
catch
(
std
::
exception
&
e
)
{
}
catch
(
std
::
exception
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
KNOWHERE_ERROR
;
}
...
...
@@ -211,22 +208,18 @@ BFIndex::Build(const Config &cfg) {
}
Status
BFIndex
::
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
{
BFIndex
::
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
{
try
{
dim
=
cfg
->
d
;
auto
dataset
=
GenDatasetWithIds
(
nb
,
dim
,
xb
,
ids
);
std
::
static_pointer_cast
<
knowhere
::
IDMAP
>
(
index_
)
->
Train
(
cfg
);
index_
->
Add
(
dataset
,
cfg
);
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_UNEXPECTED_ERROR
,
e
.
what
());
}
catch
(
std
::
exception
&
e
)
{
}
catch
(
std
::
exception
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_ERROR
,
e
.
what
());
}
...
...
@@ -235,12 +228,8 @@ BFIndex::BuildAll(const int64_t &nb,
// TODO(linxj): add lock here.
Status
IVFMixIndex
::
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
{
IVFMixIndex
::
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
{
try
{
dim
=
cfg
->
d
;
auto
dataset
=
GenDatasetWithIds
(
nb
,
dim
,
xb
,
ids
);
...
...
@@ -259,10 +248,10 @@ IVFMixIndex::BuildAll(const int64_t &nb,
WRAPPER_LOG_ERROR
<<
"Build IVFMIXIndex Failed"
;
return
Status
(
KNOWHERE_ERROR
,
"Build IVFMIXIndex Failed"
);
}
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_UNEXPECTED_ERROR
,
e
.
what
());
}
catch
(
std
::
exception
&
e
)
{
}
catch
(
std
::
exception
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_ERROR
,
e
.
what
());
}
...
...
@@ -270,12 +259,12 @@ IVFMixIndex::BuildAll(const int64_t &nb,
}
Status
IVFMixIndex
::
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
{
IVFMixIndex
::
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
{
index_
->
Load
(
index_binary
);
dim
=
Dimension
();
return
Status
::
OK
();
}
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/VecImpl.h
浏览文件 @
7633437e
...
...
@@ -15,14 +15,13 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
#include "knowhere/index/vector_index/VectorIndex.h"
#include "VecIndex.h"
#include "knowhere/index/vector_index/VectorIndex.h"
#include <utility>
#include <memory>
#include <utility>
namespace
zilliz
{
namespace
milvus
{
...
...
@@ -30,23 +29,19 @@ namespace engine {
class
VecIndexImpl
:
public
VecIndex
{
public:
explicit
VecIndexImpl
(
std
::
shared_ptr
<
zilliz
::
knowhere
::
VectorIndex
>
index
,
const
IndexType
&
type
)
explicit
VecIndexImpl
(
std
::
shared_ptr
<
zilliz
::
knowhere
::
VectorIndex
>
index
,
const
IndexType
&
type
)
:
index_
(
std
::
move
(
index
)),
type
(
type
)
{
}
Status
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
override
;
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
override
;
VecIndexPtr
CopyToGpu
(
const
int64_t
&
device_id
,
const
Config
&
cfg
)
override
;
CopyToGpu
(
const
int64_t
&
device_id
,
const
Config
&
cfg
)
override
;
VecIndexPtr
CopyToCpu
(
const
Config
&
cfg
)
override
;
CopyToCpu
(
const
Config
&
cfg
)
override
;
IndexType
GetType
()
override
;
...
...
@@ -58,13 +53,13 @@ class VecIndexImpl : public VecIndex {
Count
()
override
;
Status
Add
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
)
override
;
Add
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
)
override
;
zilliz
::
knowhere
::
BinarySet
Serialize
()
override
;
Status
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
override
;
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
override
;
VecIndexPtr
Clone
()
override
;
...
...
@@ -73,7 +68,7 @@ class VecIndexImpl : public VecIndex {
GetDeviceId
()
override
;
Status
Search
(
const
int64_t
&
nq
,
const
float
*
xq
,
float
*
dist
,
int64_t
*
ids
,
const
Config
&
cfg
)
override
;
Search
(
const
int64_t
&
nq
,
const
float
*
xq
,
float
*
dist
,
int64_t
*
ids
,
const
Config
&
cfg
)
override
;
protected:
int64_t
dim
=
0
;
...
...
@@ -85,46 +80,38 @@ class VecIndexImpl : public VecIndex {
class
IVFMixIndex
:
public
VecIndexImpl
{
public:
explicit
IVFMixIndex
(
std
::
shared_ptr
<
zilliz
::
knowhere
::
VectorIndex
>
index
,
const
IndexType
&
type
)
explicit
IVFMixIndex
(
std
::
shared_ptr
<
zilliz
::
knowhere
::
VectorIndex
>
index
,
const
IndexType
&
type
)
:
VecIndexImpl
(
std
::
move
(
index
),
type
)
{
}
Status
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
override
;
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
override
;
Status
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
override
;
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
override
;
};
class
BFIndex
:
public
VecIndexImpl
{
public:
explicit
BFIndex
(
std
::
shared_ptr
<
zilliz
::
knowhere
::
VectorIndex
>
index
)
:
VecIndexImpl
(
std
::
move
(
index
),
IndexType
::
FAISS_IDMAP
)
{
explicit
BFIndex
(
std
::
shared_ptr
<
zilliz
::
knowhere
::
VectorIndex
>
index
)
:
VecIndexImpl
(
std
::
move
(
index
),
IndexType
::
FAISS_IDMAP
)
{
}
ErrorCode
Build
(
const
Config
&
cfg
);
Build
(
const
Config
&
cfg
);
float
*
float
*
GetRawVectors
();
Status
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
override
;
int64_t
*
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
,
const
float
*
xt
)
override
;
int64_t
*
GetRawIds
();
};
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/VecIndex.cpp
浏览文件 @
7633437e
...
...
@@ -16,17 +16,17 @@
// under the License.
#include "wrapper/VecIndex.h"
#include "knowhere/index/vector_index/IndexIVF.h"
#include "VecImpl.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexIVFSQ.h"
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVFPQ.h"
#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/IndexIVF.h"
#include "knowhere/index/vector_index/IndexIVFPQ.h"
#include "knowhere/index/vector_index/IndexIVFSQ.h"
#include "knowhere/index/vector_index/IndexKDT.h"
#include "knowhere/index/vector_index/IndexNSG.h"
#include "knowhere/common/Exception.h"
#include "VecImpl.h"
#include "utils/Log.h"
#include <cuda.h>
...
...
@@ -39,18 +39,18 @@ struct FileIOReader {
std
::
fstream
fs
;
std
::
string
name
;
explicit
FileIOReader
(
const
std
::
string
&
fname
);
explicit
FileIOReader
(
const
std
::
string
&
fname
);
~
FileIOReader
();
size_t
operator
()(
void
*
ptr
,
size_t
size
);
operator
()(
void
*
ptr
,
size_t
size
);
size_t
operator
()(
void
*
ptr
,
size_t
size
,
size_t
pos
);
operator
()(
void
*
ptr
,
size_t
size
,
size_t
pos
);
};
FileIOReader
::
FileIOReader
(
const
std
::
string
&
fname
)
{
FileIOReader
::
FileIOReader
(
const
std
::
string
&
fname
)
{
name
=
fname
;
fs
=
std
::
fstream
(
name
,
std
::
ios
::
in
|
std
::
ios
::
binary
);
}
...
...
@@ -60,12 +60,12 @@ FileIOReader::~FileIOReader() {
}
size_t
FileIOReader
::
operator
()(
void
*
ptr
,
size_t
size
)
{
fs
.
read
(
reinterpret_cast
<
char
*>
(
ptr
),
size
);
FileIOReader
::
operator
()(
void
*
ptr
,
size_t
size
)
{
fs
.
read
(
reinterpret_cast
<
char
*>
(
ptr
),
size
);
}
size_t
FileIOReader
::
operator
()(
void
*
ptr
,
size_t
size
,
size_t
pos
)
{
FileIOReader
::
operator
()(
void
*
ptr
,
size_t
size
,
size_t
pos
)
{
return
0
;
}
...
...
@@ -73,12 +73,13 @@ struct FileIOWriter {
std
::
fstream
fs
;
std
::
string
name
;
explicit
FileIOWriter
(
const
std
::
string
&
fname
);
explicit
FileIOWriter
(
const
std
::
string
&
fname
);
~
FileIOWriter
();
size_t
operator
()(
void
*
ptr
,
size_t
size
);
size_t
operator
()(
void
*
ptr
,
size_t
size
);
};
FileIOWriter
::
FileIOWriter
(
const
std
::
string
&
fname
)
{
FileIOWriter
::
FileIOWriter
(
const
std
::
string
&
fname
)
{
name
=
fname
;
fs
=
std
::
fstream
(
name
,
std
::
ios
::
out
|
std
::
ios
::
binary
);
}
...
...
@@ -88,14 +89,14 @@ FileIOWriter::~FileIOWriter() {
}
size_t
FileIOWriter
::
operator
()(
void
*
ptr
,
size_t
size
)
{
fs
.
write
(
reinterpret_cast
<
char
*>
(
ptr
),
size
);
FileIOWriter
::
operator
()(
void
*
ptr
,
size_t
size
)
{
fs
.
write
(
reinterpret_cast
<
char
*>
(
ptr
),
size
);
}
VecIndexPtr
GetVecIndexFactory
(
const
IndexType
&
type
,
const
Config
&
cfg
)
{
GetVecIndexFactory
(
const
IndexType
&
type
,
const
Config
&
cfg
)
{
std
::
shared_ptr
<
zilliz
::
knowhere
::
VectorIndex
>
index
;
auto
gpu_device
=
-
1
;
// TODO(linxj): remove hardcode here
auto
gpu_device
=
-
1
;
// TODO(linxj): remove hardcode here
switch
(
type
)
{
case
IndexType
::
FAISS_IDMAP
:
{
index
=
std
::
make_shared
<
zilliz
::
knowhere
::
IDMAP
>
();
...
...
@@ -141,22 +142,20 @@ GetVecIndexFactory(const IndexType &type, const Config &cfg) {
index
=
std
::
make_shared
<
zilliz
::
knowhere
::
NSG
>
(
gpu_device
);
break
;
}
default:
{
return
nullptr
;
}
default:
{
return
nullptr
;
}
}
return
std
::
make_shared
<
VecIndexImpl
>
(
index
,
type
);
}
VecIndexPtr
LoadVecIndex
(
const
IndexType
&
index_type
,
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
{
LoadVecIndex
(
const
IndexType
&
index_type
,
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
{
auto
index
=
GetVecIndexFactory
(
index_type
);
index
->
Load
(
index_binary
);
return
index
;
}
VecIndexPtr
read_index
(
const
std
::
string
&
location
)
{
read_index
(
const
std
::
string
&
location
)
{
knowhere
::
BinarySet
load_data_list
;
FileIOReader
reader
(
location
);
reader
.
fs
.
seekg
(
0
,
reader
.
fs
.
end
);
...
...
@@ -201,28 +200,28 @@ read_index(const std::string &location) {
}
Status
write_index
(
VecIndexPtr
index
,
const
std
::
string
&
location
)
{
write_index
(
VecIndexPtr
index
,
const
std
::
string
&
location
)
{
try
{
auto
binaryset
=
index
->
Serialize
();
auto
index_type
=
index
->
GetType
();
FileIOWriter
writer
(
location
);
writer
(
&
index_type
,
sizeof
(
IndexType
));
for
(
auto
&
iter
:
binaryset
.
binary_map_
)
{
for
(
auto
&
iter
:
binaryset
.
binary_map_
)
{
auto
meta
=
iter
.
first
.
c_str
();
size_t
meta_length
=
iter
.
first
.
length
();
writer
(
&
meta_length
,
sizeof
(
meta_length
));
writer
((
void
*
)
meta
,
meta_length
);
writer
((
void
*
)
meta
,
meta_length
);
auto
binary
=
iter
.
second
;
int64_t
binary_length
=
binary
->
size
;
writer
(
&
binary_length
,
sizeof
(
binary_length
));
writer
((
void
*
)
binary
->
data
.
get
(),
binary_length
);
writer
((
void
*
)
binary
->
data
.
get
(),
binary_length
);
}
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_UNEXPECTED_ERROR
,
e
.
what
());
}
catch
(
std
::
exception
&
e
)
{
}
catch
(
std
::
exception
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
std
::
string
estring
(
e
.
what
());
if
(
estring
.
find
(
"No space left on device"
)
!=
estring
.
npos
)
{
...
...
@@ -236,7 +235,7 @@ write_index(VecIndexPtr index, const std::string &location) {
}
IndexType
ConvertToCpuIndexType
(
const
IndexType
&
type
)
{
ConvertToCpuIndexType
(
const
IndexType
&
type
)
{
// TODO(linxj): add IDMAP
switch
(
type
)
{
case
IndexType
::
FAISS_IVFFLAT_GPU
:
...
...
@@ -247,14 +246,12 @@ ConvertToCpuIndexType(const IndexType &type) {
case
IndexType
::
FAISS_IVFSQ8_MIX
:
{
return
IndexType
::
FAISS_IVFSQ8_CPU
;
}
default:
{
return
type
;
}
default:
{
return
type
;
}
}
}
IndexType
ConvertToGpuIndexType
(
const
IndexType
&
type
)
{
ConvertToGpuIndexType
(
const
IndexType
&
type
)
{
switch
(
type
)
{
case
IndexType
::
FAISS_IVFFLAT_MIX
:
case
IndexType
::
FAISS_IVFFLAT_CPU
:
{
...
...
@@ -264,12 +261,10 @@ ConvertToGpuIndexType(const IndexType &type) {
case
IndexType
::
FAISS_IVFSQ8_CPU
:
{
return
IndexType
::
FAISS_IVFSQ8_GPU
;
}
default:
{
return
type
;
}
default:
{
return
type
;
}
}
}
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
cpp/src/wrapper/VecIndex.h
浏览文件 @
7633437e
...
...
@@ -15,15 +15,14 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
#include <string>
#include <memory>
#include <string>
#include "utils/Status.h"
#include "knowhere/common/Config.h"
#include "knowhere/common/BinarySet.h"
#include "knowhere/common/Config.h"
#include "utils/Status.h"
namespace
zilliz
{
namespace
milvus
{
...
...
@@ -36,7 +35,7 @@ enum class IndexType {
FAISS_IDMAP
=
1
,
FAISS_IVFFLAT_CPU
,
FAISS_IVFFLAT_GPU
,
FAISS_IVFFLAT_MIX
,
// build on gpu and search on cpu
FAISS_IVFFLAT_MIX
,
// build on gpu and search on cpu
FAISS_IVFPQ_CPU
,
FAISS_IVFPQ_GPU
,
SPTAG_KDT_RNT_CPU
,
...
...
@@ -53,32 +52,20 @@ using VecIndexPtr = std::shared_ptr<VecIndex>;
class
VecIndex
{
public:
virtual
Status
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
=
0
,
const
float
*
xt
=
nullptr
)
=
0
;
BuildAll
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
,
const
int64_t
&
nt
=
0
,
const
float
*
xt
=
nullptr
)
=
0
;
virtual
Status
Add
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
=
Config
())
=
0
;
Add
(
const
int64_t
&
nb
,
const
float
*
xb
,
const
int64_t
*
ids
,
const
Config
&
cfg
=
Config
())
=
0
;
virtual
Status
Search
(
const
int64_t
&
nq
,
const
float
*
xq
,
float
*
dist
,
int64_t
*
ids
,
const
Config
&
cfg
=
Config
())
=
0
;
Search
(
const
int64_t
&
nq
,
const
float
*
xq
,
float
*
dist
,
int64_t
*
ids
,
const
Config
&
cfg
=
Config
())
=
0
;
virtual
VecIndexPtr
CopyToGpu
(
const
int64_t
&
device_id
,
const
Config
&
cfg
=
Config
())
=
0
;
CopyToGpu
(
const
int64_t
&
device_id
,
const
Config
&
cfg
=
Config
())
=
0
;
virtual
VecIndexPtr
CopyToCpu
(
const
Config
&
cfg
=
Config
())
=
0
;
CopyToCpu
(
const
Config
&
cfg
=
Config
())
=
0
;
virtual
VecIndexPtr
Clone
()
=
0
;
...
...
@@ -99,27 +86,27 @@ class VecIndex {
Serialize
()
=
0
;
virtual
Status
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
=
0
;
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
=
0
;
};
extern
Status
write_index
(
VecIndexPtr
index
,
const
std
::
string
&
location
);
write_index
(
VecIndexPtr
index
,
const
std
::
string
&
location
);
extern
VecIndexPtr
read_index
(
const
std
::
string
&
location
);
read_index
(
const
std
::
string
&
location
);
extern
VecIndexPtr
GetVecIndexFactory
(
const
IndexType
&
type
,
const
Config
&
cfg
=
Config
());
GetVecIndexFactory
(
const
IndexType
&
type
,
const
Config
&
cfg
=
Config
());
extern
VecIndexPtr
LoadVecIndex
(
const
IndexType
&
index_type
,
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
);
LoadVecIndex
(
const
IndexType
&
index_type
,
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
);
extern
IndexType
ConvertToCpuIndexType
(
const
IndexType
&
type
);
ConvertToCpuIndexType
(
const
IndexType
&
type
);
extern
IndexType
ConvertToGpuIndexType
(
const
IndexType
&
type
);
ConvertToGpuIndexType
(
const
IndexType
&
type
);
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
}
// namespace engine
}
// namespace milvus
}
// namespace zilliz
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录