Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
milvus
milvus
提交
7a3223dc
M
milvus
项目概览
milvus
/
milvus
12 个月 前同步成功
通知
261
Star
22476
Fork
2472
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
M
milvus
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
7a3223dc
编写于
3月 08, 2021
作者:
T
ThreadDao
提交者:
yefu.chen
3月 08, 2021
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
tags collection stats
Signed-off-by:
N
ThreadDao
<
yufen.zong@zilliz.com
>
上级
3ef18e86
变更
33
展开全部
隐藏空白更改
内联
并排
Showing
33 changed file
with
446 addition
and
684 deletion
+446
-684
internal/distributed/indexnode/service.go
internal/distributed/indexnode/service.go
+18
-16
internal/distributed/indexservice/service.go
internal/distributed/indexservice/service.go
+20
-20
internal/distributed/proxynode/client/client.go
internal/distributed/proxynode/client/client.go
+15
-2
internal/distributed/querynode/service.go
internal/distributed/querynode/service.go
+10
-8
internal/indexnode/indexnode.go
internal/indexnode/indexnode.go
+8
-8
internal/indexnode/task.go
internal/indexnode/task.go
+5
-7
internal/indexnode/task_scheduler.go
internal/indexnode/task_scheduler.go
+1
-2
internal/indexservice/indexservice.go
internal/indexservice/indexservice.go
+7
-9
internal/indexservice/meta_table.go
internal/indexservice/meta_table.go
+1
-0
internal/indexservice/node_mgr.go
internal/indexservice/node_mgr.go
+1
-0
internal/indexservice/priority_queue.go
internal/indexservice/priority_queue.go
+6
-5
internal/indexservice/task.go
internal/indexservice/task.go
+4
-6
internal/proxynode/impl.go
internal/proxynode/impl.go
+35
-35
internal/proxynode/interface.go
internal/proxynode/interface.go
+0
-105
internal/proxynode/meta_cache.go
internal/proxynode/meta_cache.go
+4
-8
internal/proxynode/meta_cache_test.go
internal/proxynode/meta_cache_test.go
+3
-0
internal/proxynode/proxy_node.go
internal/proxynode/proxy_node.go
+36
-67
internal/proxynode/segment.go
internal/proxynode/segment.go
+12
-11
internal/proxynode/task.go
internal/proxynode/task.go
+86
-84
internal/proxyservice/impl.go
internal/proxyservice/impl.go
+1
-1
internal/proxyservice/interface.go
internal/proxyservice/interface.go
+0
-20
internal/proxyservice/node_info.go
internal/proxyservice/node_info.go
+17
-29
internal/querynode/index_loader.go
internal/querynode/index_loader.go
+10
-8
internal/querynode/load_service.go
internal/querynode/load_service.go
+4
-2
internal/querynode/query_node.go
internal/querynode/query_node.go
+19
-32
internal/querynode/query_node_test.go
internal/querynode/query_node_test.go
+5
-1
internal/querynode/segment_loader.go
internal/querynode/segment_loader.go
+10
-8
internal/querynode/type_def.go
internal/querynode/type_def.go
+0
-23
internal/types/types.go
internal/types/types.go
+1
-0
internal/util/typeutil/interface.go
internal/util/typeutil/interface.go
+0
-19
tests/python_test/collection/test_collection_count.py
tests/python_test/collection/test_collection_count.py
+82
-138
tests/python_test/collection/test_collection_stats.py
tests/python_test/collection/test_collection_stats.py
+24
-9
tests/python_test/utils.py
tests/python_test/utils.py
+1
-1
未找到文件。
internal/distributed/indexnode/service.go
浏览文件 @
7a3223dc
...
...
@@ -11,21 +11,23 @@ import (
"github.com/opentracing/opentracing-go"
grpcindexserviceclient
"github.com/zilliztech/milvus-distributed/internal/distributed/indexservice/client"
"github.com/zilliztech/milvus-distributed/internal/indexnode"
"github.com/zilliztech/milvus-distributed/internal/types"
"github.com/zilliztech/milvus-distributed/internal/util/funcutil"
"google.golang.org/grpc"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
"github.com/zilliztech/milvus-distributed/internal/util/funcutil"
"google.golang.org/grpc"
)
type
Server
struct
{
i
mpl
*
indexnode
.
IndexNode
i
ndexnode
*
indexnode
.
IndexNode
grpcServer
*
grpc
.
Server
grpcErrChan
chan
error
indexServiceClient
*
grpcindexserviceclient
.
Client
indexServiceClient
types
.
IndexService
loopCtx
context
.
Context
loopCancel
func
()
loopWg
sync
.
WaitGroup
...
...
@@ -105,16 +107,16 @@ func (s *Server) init() error {
if
err
!=
nil
{
return
err
}
s
.
i
mpl
.
SetIndexServiceClient
(
s
.
indexServiceClient
)
s
.
i
ndexnode
.
SetIndexServiceClient
(
s
.
indexServiceClient
)
indexnode
.
Params
.
Init
()
indexnode
.
Params
.
Port
=
Params
.
Port
indexnode
.
Params
.
IP
=
Params
.
IP
indexnode
.
Params
.
Address
=
Params
.
Address
s
.
i
mpl
.
UpdateStateCode
(
internalpb2
.
StateCode_INITIALIZING
)
s
.
i
ndexnode
.
UpdateStateCode
(
internalpb2
.
StateCode_INITIALIZING
)
err
=
s
.
i
mpl
.
Init
()
err
=
s
.
i
ndexnode
.
Init
()
if
err
!=
nil
{
return
err
}
...
...
@@ -122,7 +124,7 @@ func (s *Server) init() error {
}
func
(
s
*
Server
)
start
()
error
{
err
:=
s
.
i
mpl
.
Start
()
err
:=
s
.
i
ndexnode
.
Start
()
if
err
!=
nil
{
return
err
}
...
...
@@ -131,8 +133,8 @@ func (s *Server) start() error {
func
(
s
*
Server
)
Stop
()
error
{
s
.
loopCancel
()
if
s
.
i
mpl
!=
nil
{
s
.
i
mpl
.
Stop
()
if
s
.
i
ndexnode
!=
nil
{
s
.
i
ndexnode
.
Stop
()
}
if
s
.
grpcServer
!=
nil
{
s
.
grpcServer
.
GracefulStop
()
...
...
@@ -143,23 +145,23 @@ func (s *Server) Stop() error {
}
func
(
s
*
Server
)
BuildIndex
(
ctx
context
.
Context
,
req
*
indexpb
.
BuildIndexCmd
)
(
*
commonpb
.
Status
,
error
)
{
return
s
.
i
mpl
.
BuildIndex
(
ctx
,
req
)
return
s
.
i
ndexnode
.
BuildIndex
(
ctx
,
req
)
}
func
(
s
*
Server
)
DropIndex
(
ctx
context
.
Context
,
request
*
indexpb
.
DropIndexRequest
)
(
*
commonpb
.
Status
,
error
)
{
return
s
.
i
mpl
.
DropIndex
(
ctx
,
request
)
return
s
.
i
ndexnode
.
DropIndex
(
ctx
,
request
)
}
func
(
s
*
Server
)
GetComponentStates
(
ctx
context
.
Context
,
empty
*
commonpb
.
Empty
)
(
*
internalpb2
.
ComponentStates
,
error
)
{
return
s
.
i
mpl
.
GetComponentStates
(
ctx
)
return
s
.
i
ndexnode
.
GetComponentStates
(
ctx
)
}
func
(
s
*
Server
)
GetTimeTickChannel
(
ctx
context
.
Context
,
empty
*
commonpb
.
Empty
)
(
*
milvuspb
.
StringResponse
,
error
)
{
return
s
.
i
mpl
.
GetTimeTickChannel
(
ctx
)
return
s
.
i
ndexnode
.
GetTimeTickChannel
(
ctx
)
}
func
(
s
*
Server
)
GetStatisticsChannel
(
ctx
context
.
Context
,
empty
*
commonpb
.
Empty
)
(
*
milvuspb
.
StringResponse
,
error
)
{
return
s
.
i
mpl
.
GetStatisticsChannel
(
ctx
)
return
s
.
i
ndexnode
.
GetStatisticsChannel
(
ctx
)
}
func
NewServer
(
ctx
context
.
Context
)
(
*
Server
,
error
)
{
...
...
@@ -173,7 +175,7 @@ func NewServer(ctx context.Context) (*Server, error) {
return
&
Server
{
loopCtx
:
ctx1
,
loopCancel
:
cancel
,
i
mpl
:
node
,
i
ndexnode
:
node
,
grpcErrChan
:
make
(
chan
error
),
},
nil
}
internal/distributed/indexservice/service.go
浏览文件 @
7a3223dc
...
...
@@ -26,7 +26,7 @@ type UniqueID = typeutil.UniqueID
type
Timestamp
=
typeutil
.
Timestamp
type
Server
struct
{
i
mpl
*
indexservice
.
IndexService
i
ndexservice
*
indexservice
.
IndexService
grpcServer
*
grpc
.
Server
grpcErrChan
chan
error
...
...
@@ -56,20 +56,20 @@ func (s *Server) init() error {
s
.
loopWg
.
Add
(
1
)
go
s
.
startGrpcLoop
(
Params
.
ServicePort
)
// wait for grpc i
mpl
loop start
// wait for grpc i
ndexservice
loop start
if
err
:=
<-
s
.
grpcErrChan
;
err
!=
nil
{
return
err
}
s
.
i
mpl
.
UpdateStateCode
(
internalpb2
.
StateCode_INITIALIZING
)
s
.
i
ndexservice
.
UpdateStateCode
(
internalpb2
.
StateCode_INITIALIZING
)
if
err
:=
s
.
i
mpl
.
Init
();
err
!=
nil
{
if
err
:=
s
.
i
ndexservice
.
Init
();
err
!=
nil
{
return
err
}
return
nil
}
func
(
s
*
Server
)
start
()
error
{
if
err
:=
s
.
i
mpl
.
Start
();
err
!=
nil
{
if
err
:=
s
.
i
ndexservice
.
Start
();
err
!=
nil
{
return
err
}
log
.
Println
(
"indexService started"
)
...
...
@@ -80,8 +80,8 @@ func (s *Server) Stop() error {
if
err
:=
s
.
closer
.
Close
();
err
!=
nil
{
return
err
}
if
s
.
i
mpl
!=
nil
{
s
.
i
mpl
.
Stop
()
if
s
.
i
ndexservice
!=
nil
{
s
.
i
ndexservice
.
Stop
()
}
s
.
loopCancel
()
...
...
@@ -94,27 +94,27 @@ func (s *Server) Stop() error {
}
func
(
s
*
Server
)
RegisterNode
(
ctx
context
.
Context
,
req
*
indexpb
.
RegisterNodeRequest
)
(
*
indexpb
.
RegisterNodeResponse
,
error
)
{
return
s
.
i
mpl
.
RegisterNode
(
ctx
,
req
)
return
s
.
i
ndexservice
.
RegisterNode
(
ctx
,
req
)
}
func
(
s
*
Server
)
BuildIndex
(
ctx
context
.
Context
,
req
*
indexpb
.
BuildIndexRequest
)
(
*
indexpb
.
BuildIndexResponse
,
error
)
{
return
s
.
i
mpl
.
BuildIndex
(
ctx
,
req
)
return
s
.
i
ndexservice
.
BuildIndex
(
ctx
,
req
)
}
func
(
s
*
Server
)
GetIndexStates
(
ctx
context
.
Context
,
req
*
indexpb
.
IndexStatesRequest
)
(
*
indexpb
.
IndexStatesResponse
,
error
)
{
return
s
.
i
mpl
.
GetIndexStates
(
ctx
,
req
)
return
s
.
i
ndexservice
.
GetIndexStates
(
ctx
,
req
)
}
func
(
s
*
Server
)
DropIndex
(
ctx
context
.
Context
,
request
*
indexpb
.
DropIndexRequest
)
(
*
commonpb
.
Status
,
error
)
{
return
s
.
i
mpl
.
DropIndex
(
ctx
,
request
)
return
s
.
i
ndexservice
.
DropIndex
(
ctx
,
request
)
}
func
(
s
*
Server
)
GetIndexFilePaths
(
ctx
context
.
Context
,
req
*
indexpb
.
IndexFilePathsRequest
)
(
*
indexpb
.
IndexFilePathsResponse
,
error
)
{
return
s
.
i
mpl
.
GetIndexFilePaths
(
ctx
,
req
)
return
s
.
i
ndexservice
.
GetIndexFilePaths
(
ctx
,
req
)
}
func
(
s
*
Server
)
NotifyBuildIndex
(
ctx
context
.
Context
,
nty
*
indexpb
.
BuildIndexNotification
)
(
*
commonpb
.
Status
,
error
)
{
return
s
.
i
mpl
.
NotifyBuildIndex
(
ctx
,
nty
)
return
s
.
i
ndexservice
.
NotifyBuildIndex
(
ctx
,
nty
)
}
func
(
s
*
Server
)
startGrpcLoop
(
grpcPort
int
)
{
...
...
@@ -147,15 +147,15 @@ func (s *Server) startGrpcLoop(grpcPort int) {
}
func
(
s
*
Server
)
GetComponentStates
(
ctx
context
.
Context
,
empty
*
commonpb
.
Empty
)
(
*
internalpb2
.
ComponentStates
,
error
)
{
return
s
.
i
mpl
.
GetComponentStates
(
ctx
)
return
s
.
i
ndexservice
.
GetComponentStates
(
ctx
)
}
func
(
s
*
Server
)
GetTimeTickChannel
(
ctx
context
.
Context
,
empty
*
commonpb
.
Empty
)
(
*
milvuspb
.
StringResponse
,
error
)
{
return
s
.
i
mpl
.
GetTimeTickChannel
(
ctx
)
return
s
.
i
ndexservice
.
GetTimeTickChannel
(
ctx
)
}
func
(
s
*
Server
)
GetStatisticsChannel
(
ctx
context
.
Context
,
empty
*
commonpb
.
Empty
)
(
*
milvuspb
.
StringResponse
,
error
)
{
return
s
.
i
mpl
.
GetStatisticsChannel
(
ctx
)
return
s
.
i
ndexservice
.
GetStatisticsChannel
(
ctx
)
}
func
NewServer
(
ctx
context
.
Context
)
(
*
Server
,
error
)
{
...
...
@@ -167,10 +167,10 @@ func NewServer(ctx context.Context) (*Server, error) {
return
nil
,
err
}
s
:=
&
Server
{
loopCtx
:
ctx1
,
loopCancel
:
cancel
,
i
mpl
:
serverImp
,
grpcErrChan
:
make
(
chan
error
),
loopCtx
:
ctx1
,
loopCancel
:
cancel
,
i
ndexservice
:
serverImp
,
grpcErrChan
:
make
(
chan
error
),
}
cfg
:=
&
config
.
Configuration
{
...
...
internal/distributed/proxynode/client/client.go
浏览文件 @
7a3223dc
...
...
@@ -6,10 +6,13 @@ import (
otgrpc
"github.com/opentracing-contrib/go-grpc"
"github.com/opentracing/opentracing-go"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/proxypb"
"github.com/zilliztech/milvus-distributed/internal/util/retry"
"google.golang.org/grpc"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
"github.com/zilliztech/milvus-distributed/internal/proto/proxypb"
)
type
Client
struct
{
...
...
@@ -18,6 +21,16 @@ type Client struct {
ctx
context
.
Context
}
func
(
c
*
Client
)
GetComponentStates
(
ctx
context
.
Context
)
(
*
internalpb2
.
ComponentStates
,
error
)
{
//TODO
panic
(
"implement me"
)
}
func
(
c
*
Client
)
GetStatisticsChannel
(
ctx
context
.
Context
)
(
*
milvuspb
.
StringResponse
,
error
)
{
//TODO
panic
(
"implement me"
)
}
func
(
c
*
Client
)
Init
()
error
{
tracer
:=
opentracing
.
GlobalTracer
()
connectGrpcFunc
:=
func
()
error
{
...
...
internal/distributed/querynode/service.go
浏览文件 @
7a3223dc
...
...
@@ -9,6 +9,8 @@ import (
"sync"
"time"
"github.com/zilliztech/milvus-distributed/internal/types"
otgrpc
"github.com/opentracing-contrib/go-grpc"
"github.com/opentracing/opentracing-go"
"github.com/uber/jaeger-client-go/config"
...
...
@@ -268,20 +270,20 @@ func (s *Server) Stop() error {
return
nil
}
func
(
s
*
Server
)
SetMasterService
(
master
qn
.
MasterServiceInterfa
ce
)
error
{
return
s
.
impl
.
SetMasterService
(
master
)
func
(
s
*
Server
)
SetMasterService
(
master
Service
types
.
MasterServi
ce
)
error
{
return
s
.
impl
.
SetMasterService
(
master
Service
)
}
func
(
s
*
Server
)
SetQueryService
(
query
qn
.
QueryServiceInterfa
ce
)
error
{
return
s
.
impl
.
SetQueryService
(
query
)
func
(
s
*
Server
)
SetQueryService
(
query
Service
types
.
QueryServi
ce
)
error
{
return
s
.
impl
.
SetQueryService
(
query
Service
)
}
func
(
s
*
Server
)
SetIndexService
(
index
qn
.
IndexServiceInterfa
ce
)
error
{
return
s
.
impl
.
SetIndexService
(
index
)
func
(
s
*
Server
)
SetIndexService
(
index
Service
types
.
IndexServi
ce
)
error
{
return
s
.
impl
.
SetIndexService
(
index
Service
)
}
func
(
s
*
Server
)
SetDataService
(
data
qn
.
DataServiceInterfa
ce
)
error
{
return
s
.
impl
.
SetDataService
(
data
)
func
(
s
*
Server
)
SetDataService
(
data
Service
types
.
DataServi
ce
)
error
{
return
s
.
impl
.
SetDataService
(
data
Service
)
}
func
(
s
*
Server
)
GetTimeTickChannel
(
ctx
context
.
Context
,
in
*
commonpb
.
Empty
)
(
*
milvuspb
.
StringResponse
,
error
)
{
...
...
internal/indexnode/indexnode.go
浏览文件 @
7a3223dc
...
...
@@ -2,24 +2,24 @@ package indexnode
import
(
"context"
"errors"
"fmt"
"io"
"log"
"time"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
"errors"
"github.com/opentracing/opentracing-go"
"github.com/uber/jaeger-client-go/config"
"github.com/zilliztech/milvus-distributed/internal/kv"
miniokv
"github.com/zilliztech/milvus-distributed/internal/kv/minio"
"github.com/zilliztech/milvus-distributed/internal/types"
"github.com/zilliztech/milvus-distributed/internal/util/funcutil"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
"github.com/zilliztech/milvus-distributed/internal/util/funcutil"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
)
const
(
...
...
@@ -39,7 +39,7 @@ type IndexNode struct {
kv
kv
.
Base
serviceClient
type
util
.
IndexServiceInterfa
ce
// method factory
serviceClient
type
s
.
IndexServi
ce
// method factory
// Add callback functions at different stages
startCallbacks
[]
func
()
...
...
@@ -155,7 +155,7 @@ func (i *IndexNode) UpdateStateCode(code internalpb2.StateCode) {
i
.
stateCode
=
code
}
func
(
i
*
IndexNode
)
SetIndexServiceClient
(
serviceClient
type
util
.
IndexServiceInterfa
ce
)
{
func
(
i
*
IndexNode
)
SetIndexServiceClient
(
serviceClient
type
s
.
IndexServi
ce
)
{
i
.
serviceClient
=
serviceClient
}
...
...
internal/indexnode/task.go
浏览文件 @
7a3223dc
...
...
@@ -2,20 +2,18 @@ package indexnode
import
(
"context"
"errors"
"fmt"
"log"
"strconv"
"errors"
"github.com/zilliztech/milvus-distributed/internal/kv"
"github.com/zilliztech/milvus-distributed/internal/storage"
"github.com/zilliztech/milvus-distributed/internal/types"
"github.com/zilliztech/milvus-distributed/internal/util/funcutil"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
"github.com/zilliztech/milvus-distributed/internal/kv"
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
"github.com/zilliztech/milvus-distributed/internal/storage"
)
const
(
...
...
@@ -75,7 +73,7 @@ type IndexBuildTask struct {
kv
kv
.
Base
savePaths
[]
string
cmd
*
indexpb
.
BuildIndexCmd
serviceClient
type
util
.
IndexServiceInterfa
ce
serviceClient
type
s
.
IndexServi
ce
nodeID
UniqueID
}
...
...
internal/indexnode/task_scheduler.go
浏览文件 @
7a3223dc
...
...
@@ -3,11 +3,10 @@ package indexnode
import
(
"container/list"
"context"
"errors"
"log"
"sync"
"errors"
"github.com/opentracing/opentracing-go"
oplog
"github.com/opentracing/opentracing-go/log"
"github.com/zilliztech/milvus-distributed/internal/kv"
...
...
internal/indexservice/indexservice.go
浏览文件 @
7a3223dc
...
...
@@ -2,28 +2,26 @@ package indexservice
import
(
"context"
"errors"
"fmt"
"log"
"sync"
"time"
"github.com/zilliztech/milvus-distributed/internal/allocator"
"github.com/zilliztech/milvus-distributed/internal/tso"
"go.etcd.io/etcd/clientv3"
"errors"
"github.com/zilliztech/milvus-distributed/internal/kv"
etcdkv
"github.com/zilliztech/milvus-distributed/internal/kv/etcd"
miniokv
"github.com/zilliztech/milvus-distributed/internal/kv/minio"
"github.com/zilliztech/milvus-distributed/internal/tso"
"github.com/zilliztech/milvus-distributed/internal/util/retry"
"github.com/zilliztech/milvus-distributed/internal/util/tsoutil"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
"go.etcd.io/etcd/clientv3"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
"github.com/zilliztech/milvus-distributed/internal/util/retry"
"github.com/zilliztech/milvus-distributed/internal/util/tsoutil"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)
const
(
...
...
internal/indexservice/meta_table.go
浏览文件 @
7a3223dc
...
...
@@ -19,6 +19,7 @@ import (
"github.com/golang/protobuf/proto"
"github.com/zilliztech/milvus-distributed/internal/kv"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
)
...
...
internal/indexservice/node_mgr.go
浏览文件 @
7a3223dc
...
...
@@ -7,6 +7,7 @@ import (
"errors"
grpcindexnodeclient
"github.com/zilliztech/milvus-distributed/internal/distributed/indexnode/client"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
...
...
internal/indexservice/priority_queue.go
浏览文件 @
7a3223dc
...
...
@@ -5,12 +5,13 @@ import (
"sync"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
"github.com/zilliztech/milvus-distributed/internal/types"
)
// An Item is something we manage in a priority queue.
type
PQItem
struct
{
value
type
util
.
IndexNodeInterfac
e
// The value of the item; arbitrary.
value
type
s
.
IndexNod
e
// The value of the item; arbitrary.
key
UniqueID
addr
*
commonpb
.
Address
priority
int
// The priority of the item in the queue.
...
...
@@ -124,7 +125,7 @@ func (pq *PriorityQueue) Peek() interface{} {
//return item.value
}
func
(
pq
*
PriorityQueue
)
PeekClient
()
(
UniqueID
,
type
util
.
IndexNodeInterfac
e
)
{
func
(
pq
*
PriorityQueue
)
PeekClient
()
(
UniqueID
,
type
s
.
IndexNod
e
)
{
item
:=
pq
.
Peek
()
if
item
==
nil
{
return
UniqueID
(
-
1
),
nil
...
...
@@ -132,11 +133,11 @@ func (pq *PriorityQueue) PeekClient() (UniqueID, typeutil.IndexNodeInterface) {
return
item
.
(
*
PQItem
)
.
key
,
item
.
(
*
PQItem
)
.
value
}
func
(
pq
*
PriorityQueue
)
PeekAllClients
()
[]
type
util
.
IndexNodeInterfac
e
{
func
(
pq
*
PriorityQueue
)
PeekAllClients
()
[]
type
s
.
IndexNod
e
{
pq
.
lock
.
RLock
()
defer
pq
.
lock
.
RUnlock
()
var
ret
[]
type
util
.
IndexNodeInterfac
e
var
ret
[]
type
s
.
IndexNod
e
for
_
,
item
:=
range
pq
.
items
{
ret
=
append
(
ret
,
item
.
value
)
}
...
...
internal/indexservice/task.go
浏览文件 @
7a3223dc
...
...
@@ -2,16 +2,14 @@ package indexservice
import
(
"context"
"errors"
"log"
"github.com/zilliztech/milvus-distributed/internal/allocator"
"github.com/zilliztech/milvus-distributed/internal/kv"
"github.com/zilliztech/milvus-distributed/internal/types"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
"errors"
"github.com/zilliztech/milvus-distributed/internal/kv"
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
)
...
...
@@ -67,7 +65,7 @@ type IndexAddTask struct {
idAllocator
*
allocator
.
GlobalIDAllocator
buildQueue
TaskQueue
kv
kv
.
Base
builderClient
type
util
.
IndexNodeInterfac
e
builderClient
type
s
.
IndexNod
e
nodeClients
*
PriorityQueue
buildClientNodeID
UniqueID
}
...
...
internal/proxynode/impl.go
浏览文件 @
7a3223dc
...
...
@@ -41,8 +41,8 @@ func (node *ProxyNode) CreateCollection(ctx context.Context, request *milvuspb.C
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
CreateCollectionRequest
:
request
,
master
Client
:
node
.
masterClient
,
dataServiceClient
:
node
.
dataService
Client
,
master
Service
:
node
.
masterService
,
dataServiceClient
:
node
.
dataService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
cct
)
...
...
@@ -71,7 +71,7 @@ func (node *ProxyNode) DropCollection(ctx context.Context, request *milvuspb.Dro
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
DropCollectionRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
dct
)
...
...
@@ -100,7 +100,7 @@ func (node *ProxyNode) HasCollection(ctx context.Context, request *milvuspb.HasC
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
HasCollectionRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
hct
)
...
...
@@ -135,7 +135,7 @@ func (node *ProxyNode) LoadCollection(ctx context.Context, request *milvuspb.Loa
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
LoadCollectionRequest
:
request
,
query
serviceClient
:
node
.
queryServiceClient
,
query
Service
:
node
.
queryService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
lct
)
...
...
@@ -164,7 +164,7 @@ func (node *ProxyNode) ReleaseCollection(ctx context.Context, request *milvuspb.
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
ReleaseCollectionRequest
:
request
,
query
serviceClient
:
node
.
queryServiceClient
,
query
Service
:
node
.
queryService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
rct
)
...
...
@@ -193,7 +193,7 @@ func (node *ProxyNode) DescribeCollection(ctx context.Context, request *milvuspb
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
DescribeCollectionRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
dct
)
...
...
@@ -225,7 +225,7 @@ func (node *ProxyNode) GetCollectionStatistics(ctx context.Context, request *mil
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
CollectionStatsRequest
:
request
,
dataService
Client
:
node
.
dataServiceClient
,
dataService
:
node
.
dataService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
g
)
...
...
@@ -257,7 +257,7 @@ func (node *ProxyNode) ShowCollections(ctx context.Context, request *milvuspb.Sh
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
ShowCollectionRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
sct
)
...
...
@@ -289,7 +289,7 @@ func (node *ProxyNode) CreatePartition(ctx context.Context, request *milvuspb.Cr
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
CreatePartitionRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
result
:
nil
,
}
...
...
@@ -316,7 +316,7 @@ func (node *ProxyNode) DropPartition(ctx context.Context, request *milvuspb.Drop
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
DropPartitionRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
result
:
nil
,
}
...
...
@@ -344,7 +344,7 @@ func (node *ProxyNode) HasPartition(ctx context.Context, request *milvuspb.HasPa
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
HasPartitionRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
result
:
nil
,
}
...
...
@@ -379,7 +379,7 @@ func (node *ProxyNode) LoadPartitions(ctx context.Context, request *milvuspb.Loa
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
LoadPartitonRequest
:
request
,
query
serviceClient
:
node
.
queryServiceClient
,
query
Service
:
node
.
queryService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
lpt
)
...
...
@@ -408,7 +408,7 @@ func (node *ProxyNode) ReleasePartitions(ctx context.Context, request *milvuspb.
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
ReleasePartitionRequest
:
request
,
query
serviceClient
:
node
.
queryServiceClient
,
query
Service
:
node
.
queryService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
rpt
)
...
...
@@ -440,7 +440,7 @@ func (node *ProxyNode) ShowPartitions(ctx context.Context, request *milvuspb.Sho
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
ShowPartitionRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
result
:
nil
,
}
...
...
@@ -473,7 +473,7 @@ func (node *ProxyNode) CreateIndex(ctx context.Context, request *milvuspb.Create
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
CreateIndexRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
cit
)
...
...
@@ -501,7 +501,7 @@ func (node *ProxyNode) DescribeIndex(ctx context.Context, request *milvuspb.Desc
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
DescribeIndexRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
dit
)
...
...
@@ -533,7 +533,7 @@ func (node *ProxyNode) DropIndex(ctx context.Context, request *milvuspb.DropInde
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
DropIndexRequest
:
request
,
master
Client
:
node
.
masterClient
,
master
Service
:
node
.
masterService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
dit
)
if
err
!=
nil
{
...
...
@@ -555,11 +555,11 @@ func (node *ProxyNode) DropIndex(ctx context.Context, request *milvuspb.DropInde
func
(
node
*
ProxyNode
)
GetIndexState
(
ctx
context
.
Context
,
request
*
milvuspb
.
IndexStateRequest
)
(
*
milvuspb
.
IndexStateResponse
,
error
)
{
// log.Println("Describe index progress for: ", request)
dipt
:=
&
GetIndexStateTask
{
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
IndexStateRequest
:
request
,
indexService
Client
:
node
.
indexServiceClient
,
master
Client
:
node
.
masterClient
,
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
IndexStateRequest
:
request
,
indexService
:
node
.
indexService
,
master
Service
:
node
.
masterService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
dipt
)
...
...
@@ -587,9 +587,9 @@ func (node *ProxyNode) GetIndexState(ctx context.Context, request *milvuspb.Inde
func
(
node
*
ProxyNode
)
Insert
(
ctx
context
.
Context
,
request
*
milvuspb
.
InsertRequest
)
(
*
milvuspb
.
InsertResponse
,
error
)
{
it
:=
&
InsertTask
{
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
dataService
Client
:
node
.
dataServiceClient
,
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
dataService
:
node
.
dataService
,
BaseInsertTask
:
BaseInsertTask
{
BaseMsg
:
msgstream
.
BaseMsg
{
HashValues
:
request
.
HashKeys
,
...
...
@@ -676,10 +676,10 @@ func (node *ProxyNode) Search(ctx context.Context, request *milvuspb.SearchReque
func
(
node
*
ProxyNode
)
Flush
(
ctx
context
.
Context
,
request
*
milvuspb
.
FlushRequest
)
(
*
commonpb
.
Status
,
error
)
{
log
.
Println
(
"AA Flush collections: "
,
request
.
CollectionNames
)
ft
:=
&
FlushTask
{
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
FlushRequest
:
request
,
dataService
Client
:
node
.
dataServiceClient
,
ctx
:
ctx
,
Condition
:
NewTaskCondition
(
ctx
),
FlushRequest
:
request
,
dataService
:
node
.
dataService
,
}
err
:=
node
.
sched
.
DdQueue
.
Enqueue
(
ft
)
...
...
@@ -716,7 +716,7 @@ func (node *ProxyNode) GetPersistentSegmentInfo(ctx context.Context, req *milvus
resp
.
Status
.
Reason
=
err
.
Error
()
return
resp
,
nil
}
infoResp
,
err
:=
node
.
dataService
Client
.
GetSegmentInfo
(
ctx
,
&
datapb
.
SegmentInfoRequest
{
infoResp
,
err
:=
node
.
dataService
.
GetSegmentInfo
(
ctx
,
&
datapb
.
SegmentInfoRequest
{
Base
:
&
commonpb
.
MsgBase
{
MsgType
:
commonpb
.
MsgType_kSegmentInfo
,
MsgID
:
0
,
...
...
@@ -763,7 +763,7 @@ func (node *ProxyNode) GetQuerySegmentInfo(ctx context.Context, req *milvuspb.Qu
resp
.
Status
.
Reason
=
err
.
Error
()
return
resp
,
nil
}
infoResp
,
err
:=
node
.
queryService
Client
.
GetSegmentInfo
(
ctx
,
&
querypb
.
SegmentInfoRequest
{
infoResp
,
err
:=
node
.
queryService
.
GetSegmentInfo
(
ctx
,
&
querypb
.
SegmentInfoRequest
{
Base
:
&
commonpb
.
MsgBase
{
MsgType
:
commonpb
.
MsgType_kSegmentInfo
,
MsgID
:
0
,
...
...
@@ -798,7 +798,7 @@ func (node *ProxyNode) GetQuerySegmentInfo(ctx context.Context, req *milvuspb.Qu
}
func
(
node
*
ProxyNode
)
getSegmentsOfCollection
(
ctx
context
.
Context
,
dbName
string
,
collectionName
string
)
([]
UniqueID
,
error
)
{
describeCollectionResponse
,
err
:=
node
.
master
Client
.
DescribeCollection
(
ctx
,
&
milvuspb
.
DescribeCollectionRequest
{
describeCollectionResponse
,
err
:=
node
.
master
Service
.
DescribeCollection
(
ctx
,
&
milvuspb
.
DescribeCollectionRequest
{
Base
:
&
commonpb
.
MsgBase
{
MsgType
:
commonpb
.
MsgType_kDescribeCollection
,
MsgID
:
0
,
...
...
@@ -815,7 +815,7 @@ func (node *ProxyNode) getSegmentsOfCollection(ctx context.Context, dbName strin
return
nil
,
errors
.
New
(
describeCollectionResponse
.
Status
.
Reason
)
}
collectionID
:=
describeCollectionResponse
.
CollectionID
showPartitionsResp
,
err
:=
node
.
master
Client
.
ShowPartitions
(
ctx
,
&
milvuspb
.
ShowPartitionRequest
{
showPartitionsResp
,
err
:=
node
.
master
Service
.
ShowPartitions
(
ctx
,
&
milvuspb
.
ShowPartitionRequest
{
Base
:
&
commonpb
.
MsgBase
{
MsgType
:
commonpb
.
MsgType_kShowPartitions
,
MsgID
:
0
,
...
...
@@ -835,7 +835,7 @@ func (node *ProxyNode) getSegmentsOfCollection(ctx context.Context, dbName strin
ret
:=
make
([]
UniqueID
,
0
)
for
_
,
partitionID
:=
range
showPartitionsResp
.
PartitionIDs
{
showSegmentResponse
,
err
:=
node
.
master
Client
.
ShowSegments
(
ctx
,
&
milvuspb
.
ShowSegmentRequest
{
showSegmentResponse
,
err
:=
node
.
master
Service
.
ShowSegments
(
ctx
,
&
milvuspb
.
ShowSegmentRequest
{
Base
:
&
commonpb
.
MsgBase
{
MsgType
:
commonpb
.
MsgType_kShowSegment
,
MsgID
:
0
,
...
...
internal/proxynode/interface.go
已删除
100644 → 0
浏览文件 @
3ef18e86
package
proxynode
import
(
"context"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/datapb"
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
"github.com/zilliztech/milvus-distributed/internal/proto/proxypb"
"github.com/zilliztech/milvus-distributed/internal/proto/querypb"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)
type
MasterClient
interface
{
CreateCollection
(
ctx
context
.
Context
,
in
*
milvuspb
.
CreateCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
DropCollection
(
ctx
context
.
Context
,
in
*
milvuspb
.
DropCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
HasCollection
(
ctx
context
.
Context
,
in
*
milvuspb
.
HasCollectionRequest
)
(
*
milvuspb
.
BoolResponse
,
error
)
DescribeCollection
(
ctx
context
.
Context
,
in
*
milvuspb
.
DescribeCollectionRequest
)
(
*
milvuspb
.
DescribeCollectionResponse
,
error
)
ShowCollections
(
ctx
context
.
Context
,
in
*
milvuspb
.
ShowCollectionRequest
)
(
*
milvuspb
.
ShowCollectionResponse
,
error
)
CreatePartition
(
ctx
context
.
Context
,
in
*
milvuspb
.
CreatePartitionRequest
)
(
*
commonpb
.
Status
,
error
)
DropPartition
(
ctx
context
.
Context
,
in
*
milvuspb
.
DropPartitionRequest
)
(
*
commonpb
.
Status
,
error
)
HasPartition
(
ctx
context
.
Context
,
in
*
milvuspb
.
HasPartitionRequest
)
(
*
milvuspb
.
BoolResponse
,
error
)
ShowPartitions
(
ctx
context
.
Context
,
in
*
milvuspb
.
ShowPartitionRequest
)
(
*
milvuspb
.
ShowPartitionResponse
,
error
)
CreateIndex
(
ctx
context
.
Context
,
in
*
milvuspb
.
CreateIndexRequest
)
(
*
commonpb
.
Status
,
error
)
DescribeIndex
(
ctx
context
.
Context
,
in
*
milvuspb
.
DescribeIndexRequest
)
(
*
milvuspb
.
DescribeIndexResponse
,
error
)
DropIndex
(
ctx
context
.
Context
,
in
*
milvuspb
.
DropIndexRequest
)
(
*
commonpb
.
Status
,
error
)
ShowSegments
(
ctx
context
.
Context
,
in
*
milvuspb
.
ShowSegmentRequest
)
(
*
milvuspb
.
ShowSegmentResponse
,
error
)
DescribeSegment
(
ctx
context
.
Context
,
in
*
milvuspb
.
DescribeSegmentRequest
)
(
*
milvuspb
.
DescribeSegmentResponse
,
error
)
}
type
IndexServiceClient
interface
{
GetIndexStates
(
ctx
context
.
Context
,
req
*
indexpb
.
IndexStatesRequest
)
(
*
indexpb
.
IndexStatesResponse
,
error
)
GetComponentStates
(
ctx
context
.
Context
)
(
*
internalpb2
.
ComponentStates
,
error
)
}
type
QueryServiceClient
interface
{
ShowCollections
(
ctx
context
.
Context
,
req
*
querypb
.
ShowCollectionRequest
)
(
*
querypb
.
ShowCollectionResponse
,
error
)
LoadCollection
(
ctx
context
.
Context
,
req
*
querypb
.
LoadCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
ReleaseCollection
(
ctx
context
.
Context
,
req
*
querypb
.
ReleaseCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
ShowPartitions
(
ctx
context
.
Context
,
req
*
querypb
.
ShowPartitionRequest
)
(
*
querypb
.
ShowPartitionResponse
,
error
)
LoadPartitions
(
ctx
context
.
Context
,
req
*
querypb
.
LoadPartitionRequest
)
(
*
commonpb
.
Status
,
error
)
ReleasePartitions
(
ctx
context
.
Context
,
req
*
querypb
.
ReleasePartitionRequest
)
(
*
commonpb
.
Status
,
error
)
CreateQueryChannel
(
ctx
context
.
Context
)
(
*
querypb
.
CreateQueryChannelResponse
,
error
)
GetPartitionStates
(
ctx
context
.
Context
,
req
*
querypb
.
PartitionStatesRequest
)
(
*
querypb
.
PartitionStatesResponse
,
error
)
//GetSearchChannelNames() ([]string, error)
//GetSearchResultChannels() ([]string, error)
GetComponentStates
(
ctx
context
.
Context
)
(
*
internalpb2
.
ComponentStates
,
error
)
GetSegmentInfo
(
ctx
context
.
Context
,
req
*
querypb
.
SegmentInfoRequest
)
(
*
querypb
.
SegmentInfoResponse
,
error
)
}
type
DataServiceClient
interface
{
AssignSegmentID
(
ctx
context
.
Context
,
req
*
datapb
.
AssignSegIDRequest
)
(
*
datapb
.
AssignSegIDResponse
,
error
)
GetInsertChannels
(
ctx
context
.
Context
,
req
*
datapb
.
InsertChannelRequest
)
(
*
internalpb2
.
StringList
,
error
)
Flush
(
ctx
context
.
Context
,
req
*
datapb
.
FlushRequest
)
(
*
commonpb
.
Status
,
error
)
GetCollectionStatistics
(
ctx
context
.
Context
,
req
*
datapb
.
CollectionStatsRequest
)
(
*
datapb
.
CollectionStatsResponse
,
error
)
GetComponentStates
(
ctx
context
.
Context
)
(
*
internalpb2
.
ComponentStates
,
error
)
GetSegmentInfo
(
ctx
context
.
Context
,
req
*
datapb
.
SegmentInfoRequest
)
(
*
datapb
.
SegmentInfoResponse
,
error
)
}
type
ProxyServiceClient
interface
{
GetTimeTickChannel
(
ctx
context
.
Context
)
(
*
milvuspb
.
StringResponse
,
error
)
RegisterNode
(
ctx
context
.
Context
,
request
*
proxypb
.
RegisterNodeRequest
)
(
*
proxypb
.
RegisterNodeResponse
,
error
)
GetComponentStates
(
ctx
context
.
Context
)
(
*
internalpb2
.
ComponentStates
,
error
)
}
type
Service
interface
{
typeutil
.
Service
InvalidateCollectionMetaCache
(
ctx
context
.
Context
,
request
*
proxypb
.
InvalidateCollMetaCacheRequest
)
(
*
commonpb
.
Status
,
error
)
CreateCollection
(
ctx
context
.
Context
,
request
*
milvuspb
.
CreateCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
DropCollection
(
ctx
context
.
Context
,
request
*
milvuspb
.
DropCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
HasCollection
(
ctx
context
.
Context
,
request
*
milvuspb
.
HasCollectionRequest
)
(
*
milvuspb
.
BoolResponse
,
error
)
LoadCollection
(
ctx
context
.
Context
,
request
*
milvuspb
.
LoadCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
ReleaseCollection
(
ctx
context
.
Context
,
request
*
milvuspb
.
ReleaseCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
DescribeCollection
(
ctx
context
.
Context
,
request
*
milvuspb
.
DescribeCollectionRequest
)
(
*
milvuspb
.
DescribeCollectionResponse
,
error
)
GetCollectionStatistics
(
ctx
context
.
Context
,
request
*
milvuspb
.
CollectionStatsRequest
)
(
*
milvuspb
.
CollectionStatsResponse
,
error
)
ShowCollections
(
ctx
context
.
Context
,
request
*
milvuspb
.
ShowCollectionRequest
)
(
*
milvuspb
.
ShowCollectionResponse
,
error
)
CreatePartition
(
ctx
context
.
Context
,
request
*
milvuspb
.
CreatePartitionRequest
)
(
*
commonpb
.
Status
,
error
)
DropPartition
(
ctx
context
.
Context
,
request
*
milvuspb
.
DropPartitionRequest
)
(
*
commonpb
.
Status
,
error
)
HasPartition
(
ctx
context
.
Context
,
request
*
milvuspb
.
HasPartitionRequest
)
(
*
milvuspb
.
BoolResponse
,
error
)
LoadPartitions
(
ctx
context
.
Context
,
request
*
milvuspb
.
LoadPartitonRequest
)
(
*
commonpb
.
Status
,
error
)
ReleasePartitions
(
ctx
context
.
Context
,
request
*
milvuspb
.
ReleasePartitionRequest
)
(
*
commonpb
.
Status
,
error
)
GetPartitionStatistics
(
ctx
context
.
Context
,
request
*
milvuspb
.
PartitionStatsRequest
)
(
*
milvuspb
.
PartitionStatsResponse
,
error
)
ShowPartitions
(
ctx
context
.
Context
,
request
*
milvuspb
.
ShowPartitionRequest
)
(
*
milvuspb
.
ShowPartitionResponse
,
error
)
CreateIndex
(
ctx
context
.
Context
,
request
*
milvuspb
.
CreateIndexRequest
)
(
*
commonpb
.
Status
,
error
)
DescribeIndex
(
ctx
context
.
Context
,
request
*
milvuspb
.
DescribeIndexRequest
)
(
*
milvuspb
.
DescribeIndexResponse
,
error
)
GetIndexState
(
ctx
context
.
Context
,
request
*
milvuspb
.
IndexStateRequest
)
(
*
milvuspb
.
IndexStateResponse
,
error
)
DropIndex
(
ctx
context
.
Context
,
request
*
milvuspb
.
DropIndexRequest
)
(
*
commonpb
.
Status
,
error
)
Insert
(
ctx
context
.
Context
,
request
*
milvuspb
.
InsertRequest
)
(
*
milvuspb
.
InsertResponse
,
error
)
Search
(
ctx
context
.
Context
,
request
*
milvuspb
.
SearchRequest
)
(
*
milvuspb
.
SearchResults
,
error
)
Flush
(
ctx
context
.
Context
,
request
*
milvuspb
.
FlushRequest
)
(
*
commonpb
.
Status
,
error
)
GetDdChannel
(
ctx
context
.
Context
,
request
*
commonpb
.
Empty
)
(
*
milvuspb
.
StringResponse
,
error
)
GetQuerySegmentInfo
(
ctx
context
.
Context
,
req
*
milvuspb
.
QuerySegmentInfoRequest
)
(
*
milvuspb
.
QuerySegmentInfoResponse
,
error
)
GetPersistentSegmentInfo
(
ctx
context
.
Context
,
req
*
milvuspb
.
PersistentSegmentInfoRequest
)
(
*
milvuspb
.
PersistentSegmentInfoResponse
,
error
)
}
internal/proxynode/meta_cache.go
浏览文件 @
7a3223dc
...
...
@@ -9,14 +9,10 @@ import (
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
"github.com/zilliztech/milvus-distributed/internal/proto/schemapb"
"github.com/zilliztech/milvus-distributed/internal/types"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)
type
MasterClientInterface
interface
{
DescribeCollection
(
ctx
context
.
Context
,
in
*
milvuspb
.
DescribeCollectionRequest
)
(
*
milvuspb
.
DescribeCollectionResponse
,
error
)
ShowPartitions
(
ctx
context
.
Context
,
in
*
milvuspb
.
ShowPartitionRequest
)
(
*
milvuspb
.
ShowPartitionResponse
,
error
)
}
type
Cache
interface
{
GetCollectionID
(
ctx
context
.
Context
,
collectionName
string
)
(
typeutil
.
UniqueID
,
error
)
GetPartitionID
(
ctx
context
.
Context
,
collectionName
string
,
partitionName
string
)
(
typeutil
.
UniqueID
,
error
)
...
...
@@ -32,7 +28,7 @@ type collectionInfo struct {
}
type
MetaCache
struct
{
client
MasterClientInterfa
ce
client
types
.
MasterServi
ce
collInfo
map
[
string
]
*
collectionInfo
mu
sync
.
RWMutex
...
...
@@ -40,7 +36,7 @@ type MetaCache struct {
var
globalMetaCache
Cache
func
InitMetaCache
(
client
MasterClientInterfa
ce
)
error
{
func
InitMetaCache
(
client
types
.
MasterServi
ce
)
error
{
var
err
error
globalMetaCache
,
err
=
NewMetaCache
(
client
)
if
err
!=
nil
{
...
...
@@ -49,7 +45,7 @@ func InitMetaCache(client MasterClientInterface) error {
return
nil
}
func
NewMetaCache
(
client
MasterClientInterfa
ce
)
(
*
MetaCache
,
error
)
{
func
NewMetaCache
(
client
types
.
MasterServi
ce
)
(
*
MetaCache
,
error
)
{
return
&
MetaCache
{
client
:
client
,
collInfo
:
map
[
string
]
*
collectionInfo
{},
...
...
internal/proxynode/meta_cache_test.go
浏览文件 @
7a3223dc
...
...
@@ -4,6 +4,8 @@ import (
"context"
"testing"
"github.com/zilliztech/milvus-distributed/internal/types"
"github.com/stretchr/testify/assert"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
...
...
@@ -12,6 +14,7 @@ import (
)
type
MockMasterClientInterface
struct
{
types
.
MasterService
}
func
(
m
*
MockMasterClientInterface
)
ShowPartitions
(
ctx
context
.
Context
,
in
*
milvuspb
.
ShowPartitionRequest
)
(
*
milvuspb
.
ShowPartitionResponse
,
error
)
{
...
...
internal/proxynode/proxy_node.go
浏览文件 @
7a3223dc
...
...
@@ -2,25 +2,22 @@ package proxynode
import
(
"context"
"
fmt
"
"
errors
"
"log"
"math/rand"
"sync"
"sync/atomic"
"time"
"github.com/zilliztech/milvus-distributed/internal/proto/proxypb"
"github.com/zilliztech/milvus-distributed/internal/util/retry"
"errors"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
"github.com/zilliztech/milvus-distributed/internal/allocator"
"github.com/zilliztech/milvus-distributed/internal/msgstream"
"github.com/zilliztech/milvus-distributed/internal/types"
"github.com/zilliztech/milvus-distributed/internal/util/funcutil"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
"github.com/zilliztech/milvus-distributed/internal/proto/proxypb"
)
type
UniqueID
=
typeutil
.
UniqueID
...
...
@@ -37,11 +34,11 @@ type ProxyNode struct {
stateCode
atomic
.
Value
master
Client
MasterClient
indexService
Client
IndexServiceClient
dataService
Client
DataServiceClient
proxyService
Client
ProxyServiceClient
queryService
Client
QueryServiceClient
master
Service
types
.
MasterService
indexService
types
.
IndexService
dataService
types
.
DataService
proxyService
types
.
ProxyService
queryService
types
.
QueryService
sched
*
TaskScheduler
tick
*
timeTick
...
...
@@ -72,39 +69,11 @@ func NewProxyNode(ctx context.Context, factory msgstream.Factory) (*ProxyNode, e
}
type
Component
interface
{
GetComponentStates
(
ctx
context
.
Context
)
(
*
internalpb2
.
ComponentStates
,
error
)
}
func
(
node
*
ProxyNode
)
waitForServiceReady
(
ctx
context
.
Context
,
service
Component
,
serviceName
string
)
error
{
checkFunc
:=
func
()
error
{
resp
,
err
:=
service
.
GetComponentStates
(
ctx
)
if
err
!=
nil
{
return
err
}
if
resp
.
Status
.
ErrorCode
!=
commonpb
.
ErrorCode_SUCCESS
{
return
errors
.
New
(
resp
.
Status
.
Reason
)
}
if
resp
.
State
.
StateCode
!=
internalpb2
.
StateCode_HEALTHY
{
return
errors
.
New
(
""
)
}
return
nil
}
// wait for 10 seconds
err
:=
retry
.
Retry
(
200
,
time
.
Millisecond
*
200
,
checkFunc
)
if
err
!=
nil
{
errMsg
:=
fmt
.
Sprintf
(
"ProxyNode wait for %s ready failed"
,
serviceName
)
return
errors
.
New
(
errMsg
)
}
return
nil
}
func
(
node
*
ProxyNode
)
Init
()
error
{
// todo wait for proxyservice state changed to Healthy
ctx
:=
context
.
Background
()
err
:=
node
.
waitForServiceReady
(
ctx
,
node
.
proxyServiceClient
,
"ProxyService"
)
err
:=
funcutil
.
WaitForComponentHealthy
(
ctx
,
node
.
proxyService
,
"ProxyService"
,
100
,
time
.
Millisecond
*
200
)
if
err
!=
nil
{
return
err
}
...
...
@@ -117,7 +86,7 @@ func (node *ProxyNode) Init() error {
},
}
response
,
err
:=
node
.
proxyService
Client
.
RegisterNode
(
ctx
,
request
)
response
,
err
:=
node
.
proxyService
.
RegisterNode
(
ctx
,
request
)
if
err
!=
nil
{
return
err
}
...
...
@@ -131,31 +100,31 @@ func (node *ProxyNode) Init() error {
}
// wait for dataservice state changed to Healthy
if
node
.
dataService
Client
!=
nil
{
err
=
node
.
waitForServiceReady
(
ctx
,
node
.
dataServiceClient
,
"DataService"
)
if
node
.
dataService
!=
nil
{
err
:=
funcutil
.
WaitForComponentHealthy
(
ctx
,
node
.
dataService
,
"DataService"
,
100
,
time
.
Millisecond
*
200
)
if
err
!=
nil
{
return
err
}
}
// wait for query
s
ervice state changed to Healthy
if
node
.
queryService
Client
!=
nil
{
err
=
node
.
waitForServiceReady
(
ctx
,
node
.
queryServiceClient
,
"QueryService"
)
// wait for query
S
ervice state changed to Healthy
if
node
.
queryService
!=
nil
{
err
:=
funcutil
.
WaitForComponentHealthy
(
ctx
,
node
.
queryService
,
"QueryService"
,
100
,
time
.
Millisecond
*
200
)
if
err
!=
nil
{
return
err
}
}
// wait for indexservice state changed to Healthy
if
node
.
indexService
Client
!=
nil
{
err
=
node
.
waitForServiceReady
(
ctx
,
node
.
indexServiceClient
,
"IndexService"
)
if
node
.
indexService
!=
nil
{
err
:=
funcutil
.
WaitForComponentHealthy
(
ctx
,
node
.
indexService
,
"IndexService"
,
100
,
time
.
Millisecond
*
200
)
if
err
!=
nil
{
return
err
}
}
if
node
.
queryService
Client
!=
nil
{
resp
,
err
:=
node
.
queryService
Client
.
CreateQueryChannel
(
ctx
)
if
node
.
queryService
!=
nil
{
resp
,
err
:=
node
.
queryService
.
CreateQueryChannel
(
ctx
)
if
err
!=
nil
{
return
err
}
...
...
@@ -168,7 +137,7 @@ func (node *ProxyNode) Init() error {
}
// todo
//Params.InsertChannelNames, err = node.dataService
Client
.GetInsertChannels()
//Params.InsertChannelNames, err = node.dataService.GetInsertChannels()
//if err != nil {
// return err
//}
...
...
@@ -204,7 +173,7 @@ func (node *ProxyNode) Init() error {
node
.
tsoAllocator
=
tsoAllocator
node
.
tsoAllocator
.
PeerID
=
Params
.
ProxyID
segAssigner
,
err
:=
NewSegIDAssigner
(
node
.
ctx
,
node
.
dataService
Client
,
node
.
lastTick
)
segAssigner
,
err
:=
NewSegIDAssigner
(
node
.
ctx
,
node
.
dataService
,
node
.
lastTick
)
if
err
!=
nil
{
panic
(
err
)
}
...
...
@@ -232,7 +201,7 @@ func (node *ProxyNode) Init() error {
}
func
(
node
*
ProxyNode
)
Start
()
error
{
err
:=
InitMetaCache
(
node
.
master
Client
)
err
:=
InitMetaCache
(
node
.
master
Service
)
if
err
!=
nil
{
return
err
}
...
...
@@ -308,22 +277,22 @@ func (node *ProxyNode) AddCloseCallback(callbacks ...func()) {
node
.
closeCallbacks
=
append
(
node
.
closeCallbacks
,
callbacks
...
)
}
func
(
node
*
ProxyNode
)
SetMasterClient
(
cli
MasterClient
)
{
node
.
master
Client
=
cli
func
(
node
*
ProxyNode
)
SetMasterClient
(
cli
types
.
MasterService
)
{
node
.
master
Service
=
cli
}
func
(
node
*
ProxyNode
)
SetIndexServiceClient
(
cli
IndexServiceClient
)
{
node
.
indexService
Client
=
cli
func
(
node
*
ProxyNode
)
SetIndexServiceClient
(
cli
types
.
IndexService
)
{
node
.
indexService
=
cli
}
func
(
node
*
ProxyNode
)
SetDataServiceClient
(
cli
DataServiceClient
)
{
node
.
dataService
Client
=
cli
func
(
node
*
ProxyNode
)
SetDataServiceClient
(
cli
types
.
DataService
)
{
node
.
dataService
=
cli
}
func
(
node
*
ProxyNode
)
SetProxyServiceClient
(
cli
ProxyServiceClient
)
{
node
.
proxyService
Client
=
cli
func
(
node
*
ProxyNode
)
SetProxyServiceClient
(
cli
types
.
ProxyService
)
{
node
.
proxyService
=
cli
}
func
(
node
*
ProxyNode
)
SetQueryServiceClient
(
cli
QueryServiceClient
)
{
node
.
queryService
Client
=
cli
func
(
node
*
ProxyNode
)
SetQueryServiceClient
(
cli
types
.
QueryService
)
{
node
.
queryService
=
cli
}
internal/proxynode/segment.go
浏览文件 @
7a3223dc
...
...
@@ -10,10 +10,11 @@ import (
"errors"
"github.com/zilliztech/milvus-distributed/internal/allocator"
"github.com/zilliztech/milvus-distributed/internal/types"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/datapb"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)
const
(
...
...
@@ -127,21 +128,21 @@ type SegIDAssigner struct {
getTickFunc
func
()
Timestamp
PeerID
UniqueID
serviceClient
DataServiceClient
countPerRPC
uint32
dataService
types
.
DataService
countPerRPC
uint32
}
func
NewSegIDAssigner
(
ctx
context
.
Context
,
client
DataServiceClient
,
getTickFunc
func
()
Timestamp
)
(
*
SegIDAssigner
,
error
)
{
func
NewSegIDAssigner
(
ctx
context
.
Context
,
dataService
types
.
DataService
,
getTickFunc
func
()
Timestamp
)
(
*
SegIDAssigner
,
error
)
{
ctx1
,
cancel
:=
context
.
WithCancel
(
ctx
)
sa
:=
&
SegIDAssigner
{
Allocator
:
Allocator
{
Ctx
:
ctx1
,
CancelFunc
:
cancel
,
},
countPerRPC
:
SegCountPerRPC
,
serviceClient
:
client
,
assignInfos
:
make
(
map
[
UniqueID
]
*
list
.
List
),
getTickFunc
:
getTickFunc
,
countPerRPC
:
SegCountPerRPC
,
dataService
:
dataService
,
assignInfos
:
make
(
map
[
UniqueID
]
*
list
.
List
),
getTickFunc
:
getTickFunc
,
}
sa
.
TChan
=
&
allocator
.
Ticker
{
UpdateInterval
:
time
.
Second
,
...
...
@@ -154,8 +155,8 @@ func NewSegIDAssigner(ctx context.Context, client DataServiceClient, getTickFunc
return
sa
,
nil
}
func
(
sa
*
SegIDAssigner
)
SetServiceClient
(
client
DataServiceClient
)
{
sa
.
serviceClient
=
client
func
(
sa
*
SegIDAssigner
)
SetServiceClient
(
client
types
.
DataService
)
{
sa
.
dataService
=
client
}
func
(
sa
*
SegIDAssigner
)
collectExpired
()
{
...
...
@@ -288,7 +289,7 @@ func (sa *SegIDAssigner) syncSegments() bool {
}
sa
.
segReqs
=
[]
*
datapb
.
SegIDRequest
{}
resp
,
err
:=
sa
.
serviceClient
.
AssignSegmentID
(
ctx
,
req
)
resp
,
err
:=
sa
.
dataService
.
AssignSegmentID
(
ctx
,
req
)
if
err
!=
nil
{
log
.
Println
(
"GRPC AssignSegmentID Failed"
,
resp
,
err
)
...
...
internal/proxynode/task.go
浏览文件 @
7a3223dc
...
...
@@ -7,6 +7,8 @@ import (
"math"
"strconv"
"github.com/zilliztech/milvus-distributed/internal/types"
"errors"
"github.com/golang/protobuf/proto"
...
...
@@ -68,10 +70,10 @@ type BaseInsertTask = msgstream.InsertMsg
type
InsertTask
struct
{
BaseInsertTask
Condition
ctx
context
.
Context
dataService
Client
DataServiceClient
result
*
milvuspb
.
InsertResponse
rowIDAllocator
*
allocator
.
IDAllocator
ctx
context
.
Context
dataService
types
.
DataService
result
*
milvuspb
.
InsertResponse
rowIDAllocator
*
allocator
.
IDAllocator
}
func
(
it
*
InsertTask
)
Ctx
()
context
.
Context
{
...
...
@@ -197,7 +199,7 @@ func (it *InsertTask) Execute(ctx context.Context) error {
stream
,
err
:=
globalInsertChannelsMap
.
getInsertMsgStream
(
collID
)
if
err
!=
nil
{
resp
,
_
:=
it
.
dataService
Client
.
GetInsertChannels
(
ctx
,
&
datapb
.
InsertChannelRequest
{
resp
,
_
:=
it
.
dataService
.
GetInsertChannels
(
ctx
,
&
datapb
.
InsertChannelRequest
{
Base
:
&
commonpb
.
MsgBase
{
MsgType
:
commonpb
.
MsgType_kInsert
,
// todo
MsgID
:
it
.
Base
.
MsgID
,
// todo
...
...
@@ -243,8 +245,8 @@ type CreateCollectionTask struct {
Condition
*
milvuspb
.
CreateCollectionRequest
ctx
context
.
Context
master
Client
MasterClient
dataServiceClient
DataServiceClient
master
Service
types
.
MasterService
dataServiceClient
types
.
DataService
result
*
commonpb
.
Status
schema
*
schemapb
.
CollectionSchema
}
...
...
@@ -352,7 +354,7 @@ func (cct *CreateCollectionTask) PreExecute(ctx context.Context) error {
func
(
cct
*
CreateCollectionTask
)
Execute
(
ctx
context
.
Context
)
error
{
var
err
error
cct
.
result
,
err
=
cct
.
master
Client
.
CreateCollection
(
ctx
,
cct
.
CreateCollectionRequest
)
cct
.
result
,
err
=
cct
.
master
Service
.
CreateCollection
(
ctx
,
cct
.
CreateCollectionRequest
)
if
err
!=
nil
{
return
err
}
...
...
@@ -392,9 +394,9 @@ func (cct *CreateCollectionTask) PostExecute(ctx context.Context) error {
type
DropCollectionTask
struct
{
Condition
*
milvuspb
.
DropCollectionRequest
ctx
context
.
Context
master
Client
MasterClient
result
*
commonpb
.
Status
ctx
context
.
Context
master
Service
types
.
MasterService
result
*
commonpb
.
Status
}
func
(
dct
*
DropCollectionTask
)
Ctx
()
context
.
Context
{
...
...
@@ -450,7 +452,7 @@ func (dct *DropCollectionTask) Execute(ctx context.Context) error {
return
err
}
dct
.
result
,
err
=
dct
.
master
Client
.
DropCollection
(
ctx
,
dct
.
DropCollectionRequest
)
dct
.
result
,
err
=
dct
.
master
Service
.
DropCollection
(
ctx
,
dct
.
DropCollectionRequest
)
if
err
!=
nil
{
return
err
}
...
...
@@ -742,9 +744,9 @@ func (st *SearchTask) PostExecute(ctx context.Context) error {
type
HasCollectionTask
struct
{
Condition
*
milvuspb
.
HasCollectionRequest
ctx
context
.
Context
master
Client
MasterClient
result
*
milvuspb
.
BoolResponse
ctx
context
.
Context
master
Service
types
.
MasterService
result
*
milvuspb
.
BoolResponse
}
func
(
hct
*
HasCollectionTask
)
Ctx
()
context
.
Context
{
...
...
@@ -796,7 +798,7 @@ func (hct *HasCollectionTask) PreExecute(ctx context.Context) error {
func
(
hct
*
HasCollectionTask
)
Execute
(
ctx
context
.
Context
)
error
{
var
err
error
hct
.
result
,
err
=
hct
.
master
Client
.
HasCollection
(
ctx
,
hct
.
HasCollectionRequest
)
hct
.
result
,
err
=
hct
.
master
Service
.
HasCollection
(
ctx
,
hct
.
HasCollectionRequest
)
if
hct
.
result
==
nil
{
return
errors
.
New
(
"has collection resp is nil"
)
}
...
...
@@ -813,9 +815,9 @@ func (hct *HasCollectionTask) PostExecute(ctx context.Context) error {
type
DescribeCollectionTask
struct
{
Condition
*
milvuspb
.
DescribeCollectionRequest
ctx
context
.
Context
master
Client
MasterClient
result
*
milvuspb
.
DescribeCollectionResponse
ctx
context
.
Context
master
Service
types
.
MasterService
result
*
milvuspb
.
DescribeCollectionResponse
}
func
(
dct
*
DescribeCollectionTask
)
Ctx
()
context
.
Context
{
...
...
@@ -867,7 +869,7 @@ func (dct *DescribeCollectionTask) PreExecute(ctx context.Context) error {
func
(
dct
*
DescribeCollectionTask
)
Execute
(
ctx
context
.
Context
)
error
{
var
err
error
dct
.
result
,
err
=
dct
.
master
Client
.
DescribeCollection
(
ctx
,
dct
.
DescribeCollectionRequest
)
dct
.
result
,
err
=
dct
.
master
Service
.
DescribeCollection
(
ctx
,
dct
.
DescribeCollectionRequest
)
if
dct
.
result
==
nil
{
return
errors
.
New
(
"has collection resp is nil"
)
}
...
...
@@ -884,9 +886,9 @@ func (dct *DescribeCollectionTask) PostExecute(ctx context.Context) error {
type
GetCollectionsStatisticsTask
struct
{
Condition
*
milvuspb
.
CollectionStatsRequest
ctx
context
.
Context
dataService
Client
DataServiceClient
result
*
milvuspb
.
CollectionStatsResponse
ctx
context
.
Context
dataService
types
.
DataService
result
*
milvuspb
.
CollectionStatsResponse
}
func
(
g
*
GetCollectionsStatisticsTask
)
Ctx
()
context
.
Context
{
...
...
@@ -947,7 +949,7 @@ func (g *GetCollectionsStatisticsTask) Execute(ctx context.Context) error {
CollectionID
:
collID
,
}
result
,
_
:=
g
.
dataService
Client
.
GetCollectionStatistics
(
ctx
,
req
)
result
,
_
:=
g
.
dataService
.
GetCollectionStatistics
(
ctx
,
req
)
if
result
==
nil
{
return
errors
.
New
(
"get collection statistics resp is nil"
)
}
...
...
@@ -971,9 +973,9 @@ func (g *GetCollectionsStatisticsTask) PostExecute(ctx context.Context) error {
type
ShowCollectionsTask
struct
{
Condition
*
milvuspb
.
ShowCollectionRequest
ctx
context
.
Context
master
Client
MasterClient
result
*
milvuspb
.
ShowCollectionResponse
ctx
context
.
Context
master
Service
types
.
MasterService
result
*
milvuspb
.
ShowCollectionResponse
}
func
(
sct
*
ShowCollectionsTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1022,7 +1024,7 @@ func (sct *ShowCollectionsTask) PreExecute(ctx context.Context) error {
func
(
sct
*
ShowCollectionsTask
)
Execute
(
ctx
context
.
Context
)
error
{
var
err
error
sct
.
result
,
err
=
sct
.
master
Client
.
ShowCollections
(
ctx
,
sct
.
ShowCollectionRequest
)
sct
.
result
,
err
=
sct
.
master
Service
.
ShowCollections
(
ctx
,
sct
.
ShowCollectionRequest
)
if
sct
.
result
==
nil
{
return
errors
.
New
(
"get collection statistics resp is nil"
)
}
...
...
@@ -1039,9 +1041,9 @@ func (sct *ShowCollectionsTask) PostExecute(ctx context.Context) error {
type
CreatePartitionTask
struct
{
Condition
*
milvuspb
.
CreatePartitionRequest
ctx
context
.
Context
master
Client
MasterClient
result
*
commonpb
.
Status
ctx
context
.
Context
master
Service
types
.
MasterService
result
*
commonpb
.
Status
}
func
(
cpt
*
CreatePartitionTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1099,7 +1101,7 @@ func (cpt *CreatePartitionTask) PreExecute(ctx context.Context) error {
}
func
(
cpt
*
CreatePartitionTask
)
Execute
(
ctx
context
.
Context
)
(
err
error
)
{
cpt
.
result
,
err
=
cpt
.
master
Client
.
CreatePartition
(
ctx
,
cpt
.
CreatePartitionRequest
)
cpt
.
result
,
err
=
cpt
.
master
Service
.
CreatePartition
(
ctx
,
cpt
.
CreatePartitionRequest
)
if
cpt
.
result
==
nil
{
return
errors
.
New
(
"get collection statistics resp is nil"
)
}
...
...
@@ -1116,9 +1118,9 @@ func (cpt *CreatePartitionTask) PostExecute(ctx context.Context) error {
type
DropPartitionTask
struct
{
Condition
*
milvuspb
.
DropPartitionRequest
ctx
context
.
Context
master
Client
MasterClient
result
*
commonpb
.
Status
ctx
context
.
Context
master
Service
types
.
MasterService
result
*
commonpb
.
Status
}
func
(
dpt
*
DropPartitionTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1176,7 +1178,7 @@ func (dpt *DropPartitionTask) PreExecute(ctx context.Context) error {
}
func
(
dpt
*
DropPartitionTask
)
Execute
(
ctx
context
.
Context
)
(
err
error
)
{
dpt
.
result
,
err
=
dpt
.
master
Client
.
DropPartition
(
ctx
,
dpt
.
DropPartitionRequest
)
dpt
.
result
,
err
=
dpt
.
master
Service
.
DropPartition
(
ctx
,
dpt
.
DropPartitionRequest
)
if
dpt
.
result
==
nil
{
return
errors
.
New
(
"get collection statistics resp is nil"
)
}
...
...
@@ -1193,9 +1195,9 @@ func (dpt *DropPartitionTask) PostExecute(ctx context.Context) error {
type
HasPartitionTask
struct
{
Condition
*
milvuspb
.
HasPartitionRequest
ctx
context
.
Context
master
Client
MasterClient
result
*
milvuspb
.
BoolResponse
ctx
context
.
Context
master
Service
types
.
MasterService
result
*
milvuspb
.
BoolResponse
}
func
(
hpt
*
HasPartitionTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1252,7 +1254,7 @@ func (hpt *HasPartitionTask) PreExecute(ctx context.Context) error {
}
func
(
hpt
*
HasPartitionTask
)
Execute
(
ctx
context
.
Context
)
(
err
error
)
{
hpt
.
result
,
err
=
hpt
.
master
Client
.
HasPartition
(
ctx
,
hpt
.
HasPartitionRequest
)
hpt
.
result
,
err
=
hpt
.
master
Service
.
HasPartition
(
ctx
,
hpt
.
HasPartitionRequest
)
if
hpt
.
result
==
nil
{
return
errors
.
New
(
"get collection statistics resp is nil"
)
}
...
...
@@ -1269,9 +1271,9 @@ func (hpt *HasPartitionTask) PostExecute(ctx context.Context) error {
type
ShowPartitionsTask
struct
{
Condition
*
milvuspb
.
ShowPartitionRequest
ctx
context
.
Context
master
Client
MasterClient
result
*
milvuspb
.
ShowPartitionResponse
ctx
context
.
Context
master
Service
types
.
MasterService
result
*
milvuspb
.
ShowPartitionResponse
}
func
(
spt
*
ShowPartitionsTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1323,7 +1325,7 @@ func (spt *ShowPartitionsTask) PreExecute(ctx context.Context) error {
func
(
spt
*
ShowPartitionsTask
)
Execute
(
ctx
context
.
Context
)
error
{
var
err
error
spt
.
result
,
err
=
spt
.
master
Client
.
ShowPartitions
(
ctx
,
spt
.
ShowPartitionRequest
)
spt
.
result
,
err
=
spt
.
master
Service
.
ShowPartitions
(
ctx
,
spt
.
ShowPartitionRequest
)
if
spt
.
result
==
nil
{
return
errors
.
New
(
"get collection statistics resp is nil"
)
}
...
...
@@ -1340,9 +1342,9 @@ func (spt *ShowPartitionsTask) PostExecute(ctx context.Context) error {
type
CreateIndexTask
struct
{
Condition
*
milvuspb
.
CreateIndexRequest
ctx
context
.
Context
master
Client
MasterClient
result
*
commonpb
.
Status
ctx
context
.
Context
master
Service
types
.
MasterService
result
*
commonpb
.
Status
}
func
(
cit
*
CreateIndexTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1401,7 +1403,7 @@ func (cit *CreateIndexTask) PreExecute(ctx context.Context) error {
func
(
cit
*
CreateIndexTask
)
Execute
(
ctx
context
.
Context
)
error
{
var
err
error
cit
.
result
,
err
=
cit
.
master
Client
.
CreateIndex
(
ctx
,
cit
.
CreateIndexRequest
)
cit
.
result
,
err
=
cit
.
master
Service
.
CreateIndex
(
ctx
,
cit
.
CreateIndexRequest
)
if
cit
.
result
==
nil
{
return
errors
.
New
(
"get collection statistics resp is nil"
)
}
...
...
@@ -1418,9 +1420,9 @@ func (cit *CreateIndexTask) PostExecute(ctx context.Context) error {
type
DescribeIndexTask
struct
{
Condition
*
milvuspb
.
DescribeIndexRequest
ctx
context
.
Context
master
Client
MasterClient
result
*
milvuspb
.
DescribeIndexResponse
ctx
context
.
Context
master
Service
types
.
MasterService
result
*
milvuspb
.
DescribeIndexResponse
}
func
(
dit
*
DescribeIndexTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1484,7 +1486,7 @@ func (dit *DescribeIndexTask) PreExecute(ctx context.Context) error {
func
(
dit
*
DescribeIndexTask
)
Execute
(
ctx
context
.
Context
)
error
{
var
err
error
dit
.
result
,
err
=
dit
.
master
Client
.
DescribeIndex
(
ctx
,
dit
.
DescribeIndexRequest
)
dit
.
result
,
err
=
dit
.
master
Service
.
DescribeIndex
(
ctx
,
dit
.
DescribeIndexRequest
)
log
.
Println
(
"YYYYY:"
,
dit
.
result
)
if
dit
.
result
==
nil
{
return
errors
.
New
(
"get collection statistics resp is nil"
)
...
...
@@ -1503,8 +1505,8 @@ type DropIndexTask struct {
Condition
ctx
context
.
Context
*
milvuspb
.
DropIndexRequest
master
Client
MasterClient
result
*
commonpb
.
Status
master
Service
types
.
MasterService
result
*
commonpb
.
Status
}
func
(
dit
*
DropIndexTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1563,7 +1565,7 @@ func (dit *DropIndexTask) PreExecute(ctx context.Context) error {
func
(
dit
*
DropIndexTask
)
Execute
(
ctx
context
.
Context
)
error
{
var
err
error
dit
.
result
,
err
=
dit
.
master
Client
.
DropIndex
(
ctx
,
dit
.
DropIndexRequest
)
dit
.
result
,
err
=
dit
.
master
Service
.
DropIndex
(
ctx
,
dit
.
DropIndexRequest
)
if
dit
.
result
==
nil
{
return
errors
.
New
(
"drop index resp is nil"
)
}
...
...
@@ -1580,10 +1582,10 @@ func (dit *DropIndexTask) PostExecute(ctx context.Context) error {
type
GetIndexStateTask
struct
{
Condition
*
milvuspb
.
IndexStateRequest
ctx
context
.
Context
indexService
Client
IndexServiceClient
master
Client
MasterClient
result
*
milvuspb
.
IndexStateResponse
ctx
context
.
Context
indexService
types
.
IndexService
master
Service
types
.
MasterService
result
*
milvuspb
.
IndexStateResponse
}
func
(
gist
*
GetIndexStateTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1658,7 +1660,7 @@ func (gist *GetIndexStateTask) Execute(ctx context.Context) error {
CollectionName
:
collectionName
,
CollectionID
:
collectionID
,
}
partitions
,
err
:=
gist
.
master
Client
.
ShowPartitions
(
ctx
,
showPartitionRequest
)
partitions
,
err
:=
gist
.
master
Service
.
ShowPartitions
(
ctx
,
showPartitionRequest
)
if
err
!=
nil
{
return
err
}
...
...
@@ -1680,7 +1682,7 @@ func (gist *GetIndexStateTask) Execute(ctx context.Context) error {
IndexName
:
gist
.
IndexName
,
}
indexDescriptionResp
,
err2
:=
gist
.
master
Client
.
DescribeIndex
(
ctx
,
&
describeIndexReq
)
indexDescriptionResp
,
err2
:=
gist
.
master
Service
.
DescribeIndex
(
ctx
,
&
describeIndexReq
)
if
err2
!=
nil
{
return
err2
}
...
...
@@ -1710,7 +1712,7 @@ func (gist *GetIndexStateTask) Execute(ctx context.Context) error {
CollectionID
:
collectionID
,
PartitionID
:
partitionID
,
}
segments
,
err
:=
gist
.
master
Client
.
ShowSegments
(
ctx
,
showSegmentsRequest
)
segments
,
err
:=
gist
.
master
Service
.
ShowSegments
(
ctx
,
showSegmentsRequest
)
if
err
!=
nil
{
return
err
}
...
...
@@ -1735,7 +1737,7 @@ func (gist *GetIndexStateTask) Execute(ctx context.Context) error {
CollectionID
:
collectionID
,
SegmentID
:
segmentID
,
}
segmentDesc
,
err
:=
gist
.
master
Client
.
DescribeSegment
(
ctx
,
describeSegmentRequest
)
segmentDesc
,
err
:=
gist
.
master
Service
.
DescribeSegment
(
ctx
,
describeSegmentRequest
)
if
err
!=
nil
{
return
err
}
...
...
@@ -1756,7 +1758,7 @@ func (gist *GetIndexStateTask) Execute(ctx context.Context) error {
return
err
}
states
,
err
:=
gist
.
indexService
Client
.
GetIndexStates
(
ctx
,
getIndexStatesRequest
)
states
,
err
:=
gist
.
indexService
.
GetIndexStates
(
ctx
,
getIndexStatesRequest
)
if
err
!=
nil
{
return
err
}
...
...
@@ -1799,9 +1801,9 @@ func (gist *GetIndexStateTask) PostExecute(ctx context.Context) error {
type
FlushTask
struct
{
Condition
*
milvuspb
.
FlushRequest
ctx
context
.
Context
dataService
Client
DataServiceClient
result
*
commonpb
.
Status
ctx
context
.
Context
dataService
types
.
DataService
result
*
commonpb
.
Status
}
func
(
ft
*
FlushTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1864,7 +1866,7 @@ func (ft *FlushTask) Execute(ctx context.Context) error {
CollectionID
:
collID
,
}
var
status
*
commonpb
.
Status
status
,
_
=
ft
.
dataService
Client
.
Flush
(
ctx
,
flushReq
)
status
,
_
=
ft
.
dataService
.
Flush
(
ctx
,
flushReq
)
if
status
==
nil
{
return
errors
.
New
(
"flush resp is nil"
)
}
...
...
@@ -1885,9 +1887,9 @@ func (ft *FlushTask) PostExecute(ctx context.Context) error {
type
LoadCollectionTask
struct
{
Condition
*
milvuspb
.
LoadCollectionRequest
ctx
context
.
Context
query
serviceClient
QueryServiceClient
result
*
commonpb
.
Status
ctx
context
.
Context
query
Service
types
.
QueryService
result
*
commonpb
.
Status
}
func
(
lct
*
LoadCollectionTask
)
Ctx
()
context
.
Context
{
...
...
@@ -1961,7 +1963,7 @@ func (lct *LoadCollectionTask) Execute(ctx context.Context) (err error) {
CollectionID
:
collID
,
Schema
:
collSchema
,
}
lct
.
result
,
err
=
lct
.
query
serviceClient
.
LoadCollection
(
ctx
,
request
)
lct
.
result
,
err
=
lct
.
query
Service
.
LoadCollection
(
ctx
,
request
)
return
err
}
...
...
@@ -1972,9 +1974,9 @@ func (lct *LoadCollectionTask) PostExecute(ctx context.Context) error {
type
ReleaseCollectionTask
struct
{
Condition
*
milvuspb
.
ReleaseCollectionRequest
ctx
context
.
Context
query
serviceClient
QueryServiceClient
result
*
commonpb
.
Status
ctx
context
.
Context
query
Service
types
.
QueryService
result
*
commonpb
.
Status
}
func
(
rct
*
ReleaseCollectionTask
)
Ctx
()
context
.
Context
{
...
...
@@ -2042,7 +2044,7 @@ func (rct *ReleaseCollectionTask) Execute(ctx context.Context) (err error) {
DbID
:
0
,
CollectionID
:
collID
,
}
rct
.
result
,
err
=
rct
.
query
serviceClient
.
ReleaseCollection
(
ctx
,
request
)
rct
.
result
,
err
=
rct
.
query
Service
.
ReleaseCollection
(
ctx
,
request
)
return
err
}
...
...
@@ -2053,9 +2055,9 @@ func (rct *ReleaseCollectionTask) PostExecute(ctx context.Context) error {
type
LoadPartitionTask
struct
{
Condition
*
milvuspb
.
LoadPartitonRequest
ctx
context
.
Context
query
serviceClient
QueryServiceClient
result
*
commonpb
.
Status
ctx
context
.
Context
query
Service
types
.
QueryService
result
*
commonpb
.
Status
}
func
(
lpt
*
LoadPartitionTask
)
ID
()
UniqueID
{
...
...
@@ -2133,7 +2135,7 @@ func (lpt *LoadPartitionTask) Execute(ctx context.Context) error {
PartitionIDs
:
partitionIDs
,
Schema
:
collSchema
,
}
lpt
.
result
,
err
=
lpt
.
query
serviceClient
.
LoadPartitions
(
ctx
,
request
)
lpt
.
result
,
err
=
lpt
.
query
Service
.
LoadPartitions
(
ctx
,
request
)
return
err
}
...
...
@@ -2144,9 +2146,9 @@ func (lpt *LoadPartitionTask) PostExecute(ctx context.Context) error {
type
ReleasePartitionTask
struct
{
Condition
*
milvuspb
.
ReleasePartitionRequest
ctx
context
.
Context
query
serviceClient
QueryServiceClient
result
*
commonpb
.
Status
ctx
context
.
Context
query
Service
types
.
QueryService
result
*
commonpb
.
Status
}
func
(
rpt
*
ReleasePartitionTask
)
Ctx
()
context
.
Context
{
...
...
@@ -2223,7 +2225,7 @@ func (rpt *ReleasePartitionTask) Execute(ctx context.Context) (err error) {
CollectionID
:
collID
,
PartitionIDs
:
partitionIDs
,
}
rpt
.
result
,
err
=
rpt
.
query
serviceClient
.
ReleasePartitions
(
ctx
,
request
)
rpt
.
result
,
err
=
rpt
.
query
Service
.
ReleasePartitions
(
ctx
,
request
)
return
err
}
...
...
internal/proxyservice/impl.go
浏览文件 @
7a3223dc
...
...
@@ -156,7 +156,7 @@ func (s *ProxyService) Stop() error {
if
err
!=
nil
{
panic
(
err
)
}
log
.
Println
(
"stop all node
client
s ..."
)
log
.
Println
(
"stop all node
ProxyNode
s ..."
)
s
.
cancel
()
...
...
internal/proxyservice/interface.go
已删除
100644 → 0
浏览文件 @
3ef18e86
package
proxyservice
import
(
"context"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
"github.com/zilliztech/milvus-distributed/internal/proto/proxypb"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)
type
Service
interface
{
typeutil
.
Component
typeutil
.
TimeTickHandler
RegisterLink
(
ctx
context
.
Context
)
(
*
milvuspb
.
RegisterLinkResponse
,
error
)
RegisterNode
(
ctx
context
.
Context
,
request
*
proxypb
.
RegisterNodeRequest
)
(
*
proxypb
.
RegisterNodeResponse
,
error
)
// TODO: i'm sure it's not a best way to keep consistency, fix me
InvalidateCollectionMetaCache
(
ctx
context
.
Context
,
request
*
proxypb
.
InvalidateCollMetaCacheRequest
)
(
*
commonpb
.
Status
,
error
)
}
internal/proxyservice/node_info.go
浏览文件 @
7a3223dc
...
...
@@ -2,19 +2,15 @@ package proxyservice
import
(
"context"
"errors"
"log"
"math/rand"
"strconv"
"sync"
"time"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
grpcproxynodeclient
"github.com/zilliztech/milvus-distributed/internal/distributed/proxynode/client"
"errors"
"github.com/zilliztech/milvus-distributed/internal/proto/proxypb"
"github.com/zilliztech/milvus-distributed/internal/types"
)
type
NodeInfo
struct
{
...
...
@@ -22,21 +18,13 @@ type NodeInfo struct {
port
int64
}
type
NodeClient
interface
{
Init
()
error
Start
()
error
Stop
()
error
InvalidateCollectionMetaCache
(
ctx
context
.
Context
,
request
*
proxypb
.
InvalidateCollMetaCacheRequest
)
(
*
commonpb
.
Status
,
error
)
}
type
GlobalNodeInfoTable
struct
{
mtx
sync
.
RWMutex
nodeIDs
[]
UniqueID
infos
map
[
UniqueID
]
*
NodeInfo
createClientMtx
sync
.
RWMutex
// lazy creating, so len(
client
s) <= len(infos)
clients
map
[
UniqueID
]
NodeClient
// lazy creating, so len(
ProxyNode
s) <= len(infos)
ProxyNodes
map
[
UniqueID
]
types
.
ProxyNode
}
func
(
table
*
GlobalNodeInfoTable
)
randomPick
()
UniqueID
{
...
...
@@ -82,22 +70,22 @@ func (table *GlobalNodeInfoTable) Register(id UniqueID, info *NodeInfo) error {
func
(
table
*
GlobalNodeInfoTable
)
createClients
()
error
{
log
.
Println
(
"infos: "
,
table
.
infos
)
log
.
Println
(
"
clients: "
,
table
.
client
s
)
if
len
(
table
.
client
s
)
==
len
(
table
.
infos
)
{
log
.
Println
(
"
ProxyNodes: "
,
table
.
ProxyNode
s
)
if
len
(
table
.
ProxyNode
s
)
==
len
(
table
.
infos
)
{
return
nil
}
for
nodeID
,
info
:=
range
table
.
infos
{
_
,
ok
:=
table
.
client
s
[
nodeID
]
_
,
ok
:=
table
.
ProxyNode
s
[
nodeID
]
if
!
ok
{
log
.
Println
(
info
)
table
.
client
s
[
nodeID
]
=
grpcproxynodeclient
.
NewClient
(
context
.
Background
(),
info
.
ip
+
":"
+
strconv
.
Itoa
(
int
(
info
.
port
)))
table
.
ProxyNode
s
[
nodeID
]
=
grpcproxynodeclient
.
NewClient
(
context
.
Background
(),
info
.
ip
+
":"
+
strconv
.
Itoa
(
int
(
info
.
port
)))
var
err
error
err
=
table
.
client
s
[
nodeID
]
.
Init
()
err
=
table
.
ProxyNode
s
[
nodeID
]
.
Init
()
if
err
!=
nil
{
panic
(
err
)
}
err
=
table
.
client
s
[
nodeID
]
.
Start
()
err
=
table
.
ProxyNode
s
[
nodeID
]
.
Start
()
if
err
!=
nil
{
panic
(
err
)
}
...
...
@@ -116,18 +104,18 @@ func (table *GlobalNodeInfoTable) ReleaseAllClients() error {
}()
var
err
error
for
id
,
client
:=
range
table
.
client
s
{
for
id
,
client
:=
range
table
.
ProxyNode
s
{
err
=
client
.
Stop
()
if
err
!=
nil
{
panic
(
err
)
}
delete
(
table
.
client
s
,
id
)
delete
(
table
.
ProxyNode
s
,
id
)
}
return
nil
}
func
(
table
*
GlobalNodeInfoTable
)
ObtainAllClients
()
(
map
[
UniqueID
]
NodeClient
,
error
)
{
func
(
table
*
GlobalNodeInfoTable
)
ObtainAllClients
()
(
map
[
UniqueID
]
types
.
ProxyNode
,
error
)
{
table
.
mtx
.
RLock
()
defer
table
.
mtx
.
RUnlock
()
...
...
@@ -136,13 +124,13 @@ func (table *GlobalNodeInfoTable) ObtainAllClients() (map[UniqueID]NodeClient, e
err
:=
table
.
createClients
()
return
table
.
client
s
,
err
return
table
.
ProxyNode
s
,
err
}
func
NewGlobalNodeInfoTable
()
*
GlobalNodeInfoTable
{
return
&
GlobalNodeInfoTable
{
nodeIDs
:
make
([]
UniqueID
,
0
),
infos
:
make
(
map
[
UniqueID
]
*
NodeInfo
),
clients
:
make
(
map
[
UniqueID
]
NodeClient
),
nodeIDs
:
make
([]
UniqueID
,
0
),
infos
:
make
(
map
[
UniqueID
]
*
NodeInfo
),
ProxyNodes
:
make
(
map
[
UniqueID
]
types
.
ProxyNode
),
}
}
internal/querynode/index_loader.go
浏览文件 @
7a3223dc
...
...
@@ -10,6 +10,8 @@ import (
"sync"
"time"
"github.com/zilliztech/milvus-distributed/internal/types"
"errors"
"go.uber.org/zap"
...
...
@@ -31,8 +33,8 @@ type indexLoader struct {
fieldIndexes
map
[
string
][]
*
internalpb2
.
IndexStats
fieldStatsChan
chan
[]
*
internalpb2
.
FieldStats
master
Client
MasterServiceInterfa
ce
index
Client
IndexServiceInterfa
ce
master
Service
types
.
MasterServi
ce
index
Service
types
.
IndexServi
ce
kv
kv
.
Base
// minio kv
}
...
...
@@ -315,7 +317,7 @@ func (loader *indexLoader) getIndexInfo(collectionID UniqueID, segmentID UniqueI
CollectionID
:
collectionID
,
SegmentID
:
segmentID
,
}
response
,
err
:=
loader
.
master
Client
.
DescribeSegment
(
ctx
,
req
)
response
,
err
:=
loader
.
master
Service
.
DescribeSegment
(
ctx
,
req
)
if
err
!=
nil
{
return
0
,
0
,
err
}
...
...
@@ -327,14 +329,14 @@ func (loader *indexLoader) getIndexInfo(collectionID UniqueID, segmentID UniqueI
func
(
loader
*
indexLoader
)
getIndexPaths
(
indexBuildID
UniqueID
)
([]
string
,
error
)
{
ctx
:=
context
.
TODO
()
if
loader
.
index
Client
==
nil
{
if
loader
.
index
Service
==
nil
{
return
nil
,
errors
.
New
(
"null index service client"
)
}
indexFilePathRequest
:=
&
indexpb
.
IndexFilePathsRequest
{
IndexBuildIDs
:
[]
UniqueID
{
indexBuildID
},
}
pathResponse
,
err
:=
loader
.
index
Client
.
GetIndexFilePaths
(
ctx
,
indexFilePathRequest
)
pathResponse
,
err
:=
loader
.
index
Service
.
GetIndexFilePaths
(
ctx
,
indexFilePathRequest
)
if
err
!=
nil
||
pathResponse
.
Status
.
ErrorCode
!=
commonpb
.
ErrorCode_SUCCESS
{
return
nil
,
err
}
...
...
@@ -389,7 +391,7 @@ func (loader *indexLoader) loadIndexDelayed(collectionID, segmentID UniqueID, in
return
nil
}
func
newIndexLoader
(
ctx
context
.
Context
,
master
Client
MasterServiceInterface
,
indexClient
IndexServiceInterfa
ce
,
replica
ReplicaInterface
)
*
indexLoader
{
func
newIndexLoader
(
ctx
context
.
Context
,
master
Service
types
.
MasterService
,
indexService
types
.
IndexServi
ce
,
replica
ReplicaInterface
)
*
indexLoader
{
option
:=
&
minioKV
.
Option
{
Address
:
Params
.
MinioEndPoint
,
AccessKeyID
:
Params
.
MinioAccessKeyID
,
...
...
@@ -410,8 +412,8 @@ func newIndexLoader(ctx context.Context, masterClient MasterServiceInterface, in
fieldIndexes
:
make
(
map
[
string
][]
*
internalpb2
.
IndexStats
),
fieldStatsChan
:
make
(
chan
[]
*
internalpb2
.
FieldStats
,
1
),
master
Client
:
masterClient
,
index
Client
:
indexClient
,
master
Service
:
masterService
,
index
Service
:
indexService
,
kv
:
client
,
}
...
...
internal/querynode/load_service.go
浏览文件 @
7a3223dc
...
...
@@ -6,6 +6,8 @@ import (
"sync"
"time"
"github.com/zilliztech/milvus-distributed/internal/types"
"errors"
"go.uber.org/zap"
...
...
@@ -161,10 +163,10 @@ func (s *loadService) loadSegmentInternal(collectionID UniqueID, partitionID Uni
return
nil
}
func
newLoadService
(
ctx
context
.
Context
,
master
Client
MasterServiceInterface
,
dataClient
DataServiceInterface
,
indexClient
IndexServiceInterfa
ce
,
replica
ReplicaInterface
,
dmStream
msgstream
.
MsgStream
)
*
loadService
{
func
newLoadService
(
ctx
context
.
Context
,
master
Service
types
.
MasterService
,
dataService
types
.
DataService
,
indexService
types
.
IndexServi
ce
,
replica
ReplicaInterface
,
dmStream
msgstream
.
MsgStream
)
*
loadService
{
ctx1
,
cancel
:=
context
.
WithCancel
(
ctx
)
segLoader
:=
newSegmentLoader
(
ctx1
,
master
Client
,
indexClient
,
dataClient
,
replica
,
dmStream
)
segLoader
:=
newSegmentLoader
(
ctx1
,
master
Service
,
indexService
,
dataService
,
replica
,
dmStream
)
return
&
loadService
{
ctx
:
ctx1
,
...
...
internal/querynode/query_node.go
浏览文件 @
7a3223dc
...
...
@@ -18,6 +18,8 @@ import (
"strings"
"sync/atomic"
"github.com/zilliztech/milvus-distributed/internal/types"
"errors"
"go.uber.org/zap"
...
...
@@ -32,21 +34,6 @@ import (
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)
type
Node
interface
{
typeutil
.
Component
AddQueryChannel
(
ctx
context
.
Context
,
in
*
queryPb
.
AddQueryChannelsRequest
)
(
*
commonpb
.
Status
,
error
)
RemoveQueryChannel
(
ctx
context
.
Context
,
in
*
queryPb
.
RemoveQueryChannelsRequest
)
(
*
commonpb
.
Status
,
error
)
WatchDmChannels
(
ctx
context
.
Context
,
in
*
queryPb
.
WatchDmChannelsRequest
)
(
*
commonpb
.
Status
,
error
)
LoadSegments
(
ctx
context
.
Context
,
in
*
queryPb
.
LoadSegmentRequest
)
(
*
commonpb
.
Status
,
error
)
ReleaseCollection
(
ctx
context
.
Context
,
in
*
queryPb
.
ReleaseCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
ReleasePartitions
(
ctx
context
.
Context
,
in
*
queryPb
.
ReleasePartitionRequest
)
(
*
commonpb
.
Status
,
error
)
ReleaseSegments
(
ctx
context
.
Context
,
in
*
queryPb
.
ReleaseSegmentRequest
)
(
*
commonpb
.
Status
,
error
)
GetSegmentInfo
(
ctx
context
.
Context
,
in
*
queryPb
.
SegmentInfoRequest
)
(
*
queryPb
.
SegmentInfoResponse
,
error
)
}
type
QueryService
=
typeutil
.
QueryServiceInterface
type
QueryNode
struct
{
typeutil
.
Service
...
...
@@ -66,10 +53,10 @@ type QueryNode struct {
statsService
*
statsService
// clients
master
Client
MasterServiceInterfa
ce
query
Client
QueryServiceInterfa
ce
index
Client
IndexServiceInterfa
ce
data
Client
DataServiceInterfa
ce
master
Service
types
.
MasterServi
ce
query
Service
types
.
QueryServi
ce
index
Service
types
.
IndexServi
ce
data
Service
types
.
DataServi
ce
msFactory
msgstream
.
Factory
}
...
...
@@ -127,7 +114,7 @@ func (node *QueryNode) Init() error {
},
}
resp
,
err
:=
node
.
query
Client
.
RegisterNode
(
ctx
,
registerReq
)
resp
,
err
:=
node
.
query
Service
.
RegisterNode
(
ctx
,
registerReq
)
if
err
!=
nil
{
panic
(
err
)
}
...
...
@@ -152,15 +139,15 @@ func (node *QueryNode) Init() error {
log
.
Debug
(
""
,
zap
.
Int64
(
"QueryNodeID"
,
Params
.
QueryNodeID
))
if
node
.
master
Client
==
nil
{
if
node
.
master
Service
==
nil
{
log
.
Error
(
"null master service detected"
)
}
if
node
.
index
Client
==
nil
{
if
node
.
index
Service
==
nil
{
log
.
Error
(
"null index service detected"
)
}
if
node
.
data
Client
==
nil
{
if
node
.
data
Service
==
nil
{
log
.
Error
(
"null data service detected"
)
}
...
...
@@ -183,7 +170,7 @@ func (node *QueryNode) Start() error {
node
.
searchService
=
newSearchService
(
node
.
queryNodeLoopCtx
,
node
.
replica
,
node
.
msFactory
)
//node.metaService = newMetaService(node.queryNodeLoopCtx, node.replica)
node
.
loadService
=
newLoadService
(
node
.
queryNodeLoopCtx
,
node
.
master
Client
,
node
.
dataClient
,
node
.
indexClient
,
node
.
replica
,
node
.
dataSyncService
.
dmStream
)
node
.
loadService
=
newLoadService
(
node
.
queryNodeLoopCtx
,
node
.
master
Service
,
node
.
dataService
,
node
.
indexService
,
node
.
replica
,
node
.
dataSyncService
.
dmStream
)
node
.
statsService
=
newStatsService
(
node
.
queryNodeLoopCtx
,
node
.
replica
,
node
.
loadService
.
segLoader
.
indexLoader
.
fieldStatsChan
,
node
.
msFactory
)
// start services
...
...
@@ -223,35 +210,35 @@ func (node *QueryNode) UpdateStateCode(code internalpb2.StateCode) {
node
.
stateCode
.
Store
(
code
)
}
func
(
node
*
QueryNode
)
SetMasterService
(
master
MasterServiceInterfa
ce
)
error
{
func
(
node
*
QueryNode
)
SetMasterService
(
master
types
.
MasterServi
ce
)
error
{
if
master
==
nil
{
return
errors
.
New
(
"null master service interface"
)
}
node
.
master
Client
=
master
node
.
master
Service
=
master
return
nil
}
func
(
node
*
QueryNode
)
SetQueryService
(
query
QueryServiceInterfa
ce
)
error
{
func
(
node
*
QueryNode
)
SetQueryService
(
query
types
.
QueryServi
ce
)
error
{
if
query
==
nil
{
return
errors
.
New
(
"null query service interface"
)
}
node
.
query
Client
=
query
node
.
query
Service
=
query
return
nil
}
func
(
node
*
QueryNode
)
SetIndexService
(
index
IndexServiceInterfa
ce
)
error
{
func
(
node
*
QueryNode
)
SetIndexService
(
index
types
.
IndexServi
ce
)
error
{
if
index
==
nil
{
return
errors
.
New
(
"null index service interface"
)
}
node
.
index
Client
=
index
node
.
index
Service
=
index
return
nil
}
func
(
node
*
QueryNode
)
SetDataService
(
data
DataServiceInterfa
ce
)
error
{
func
(
node
*
QueryNode
)
SetDataService
(
data
types
.
DataServi
ce
)
error
{
if
data
==
nil
{
return
errors
.
New
(
"null data service interface"
)
}
node
.
data
Client
=
data
node
.
data
Service
=
data
return
nil
}
...
...
internal/querynode/query_node_test.go
浏览文件 @
7a3223dc
...
...
@@ -8,6 +8,8 @@ import (
"testing"
"time"
"github.com/zilliztech/milvus-distributed/internal/types"
"github.com/stretchr/testify/assert"
"github.com/zilliztech/milvus-distributed/internal/msgstream/pulsarms"
...
...
@@ -23,7 +25,9 @@ const debug = false
const
defaultPartitionID
=
UniqueID
(
2021
)
type
queryServiceMock
struct
{}
type
queryServiceMock
struct
{
types
.
QueryService
}
func
setup
()
{
os
.
Setenv
(
"QUERY_NODE_ID"
,
"1"
)
...
...
internal/querynode/segment_loader.go
浏览文件 @
7a3223dc
...
...
@@ -4,6 +4,8 @@ import (
"context"
"strconv"
"github.com/zilliztech/milvus-distributed/internal/types"
"errors"
"github.com/zilliztech/milvus-distributed/internal/kv"
...
...
@@ -21,7 +23,7 @@ type segmentLoader struct {
dmStream
msgstream
.
MsgStream
data
Client
DataServiceInterfa
ce
data
Service
types
.
DataServi
ce
kv
kv
.
Base
// minio kv
iCodec
*
storage
.
InsertCodec
...
...
@@ -31,7 +33,7 @@ type segmentLoader struct {
func
(
loader
*
segmentLoader
)
getInsertBinlogPaths
(
segmentID
UniqueID
)
([]
*
internalpb2
.
StringList
,
[]
int64
,
error
)
{
ctx
:=
context
.
TODO
()
if
loader
.
data
Client
==
nil
{
if
loader
.
data
Service
==
nil
{
return
nil
,
nil
,
errors
.
New
(
"null data service client"
)
}
...
...
@@ -39,7 +41,7 @@ func (loader *segmentLoader) getInsertBinlogPaths(segmentID UniqueID) ([]*intern
SegmentID
:
segmentID
,
}
pathResponse
,
err
:=
loader
.
data
Client
.
GetInsertBinlogPaths
(
ctx
,
insertBinlogPathRequest
)
pathResponse
,
err
:=
loader
.
data
Service
.
GetInsertBinlogPaths
(
ctx
,
insertBinlogPathRequest
)
if
err
!=
nil
||
pathResponse
.
Status
.
ErrorCode
!=
commonpb
.
ErrorCode_SUCCESS
{
return
nil
,
nil
,
err
}
...
...
@@ -53,14 +55,14 @@ func (loader *segmentLoader) getInsertBinlogPaths(segmentID UniqueID) ([]*intern
func
(
loader
*
segmentLoader
)
GetSegmentStates
(
segmentID
UniqueID
)
(
*
datapb
.
SegmentStatesResponse
,
error
)
{
ctx
:=
context
.
TODO
()
if
loader
.
data
Client
==
nil
{
if
loader
.
data
Service
==
nil
{
return
nil
,
errors
.
New
(
"null data service client"
)
}
segmentStatesRequest
:=
&
datapb
.
SegmentStatesRequest
{
SegmentIDs
:
[]
int64
{
segmentID
},
}
statesResponse
,
err
:=
loader
.
data
Client
.
GetSegmentStates
(
ctx
,
segmentStatesRequest
)
statesResponse
,
err
:=
loader
.
data
Service
.
GetSegmentStates
(
ctx
,
segmentStatesRequest
)
if
err
!=
nil
||
statesResponse
.
Status
.
ErrorCode
!=
commonpb
.
ErrorCode_SUCCESS
{
return
nil
,
err
}
...
...
@@ -191,7 +193,7 @@ func (loader *segmentLoader) loadSegmentFieldsData(segment *Segment, targetField
return
nil
}
func
newSegmentLoader
(
ctx
context
.
Context
,
master
Client
MasterServiceInterface
,
indexClient
IndexServiceInterface
,
dataClient
DataServiceInterfa
ce
,
replica
ReplicaInterface
,
dmStream
msgstream
.
MsgStream
)
*
segmentLoader
{
func
newSegmentLoader
(
ctx
context
.
Context
,
master
Service
types
.
MasterService
,
indexService
types
.
IndexService
,
dataService
types
.
DataServi
ce
,
replica
ReplicaInterface
,
dmStream
msgstream
.
MsgStream
)
*
segmentLoader
{
option
:=
&
minioKV
.
Option
{
Address
:
Params
.
MinioEndPoint
,
AccessKeyID
:
Params
.
MinioAccessKeyID
,
...
...
@@ -206,13 +208,13 @@ func newSegmentLoader(ctx context.Context, masterClient MasterServiceInterface,
panic
(
err
)
}
iLoader
:=
newIndexLoader
(
ctx
,
master
Client
,
indexClient
,
replica
)
iLoader
:=
newIndexLoader
(
ctx
,
master
Service
,
indexService
,
replica
)
return
&
segmentLoader
{
replica
:
replica
,
dmStream
:
dmStream
,
data
Client
:
dataClient
,
data
Service
:
dataService
,
kv
:
client
,
iCodec
:
&
storage
.
InsertCodec
{},
...
...
internal/querynode/type_def.go
浏览文件 @
7a3223dc
package
querynode
import
(
"context"
"github.com/zilliztech/milvus-distributed/internal/proto/datapb"
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
"github.com/zilliztech/milvus-distributed/internal/proto/querypb"
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)
...
...
@@ -22,20 +16,3 @@ type TimeRange struct {
timestampMin
Timestamp
timestampMax
Timestamp
}
type
MasterServiceInterface
interface
{
DescribeSegment
(
ctx
context
.
Context
,
in
*
milvuspb
.
DescribeSegmentRequest
)
(
*
milvuspb
.
DescribeSegmentResponse
,
error
)
}
type
QueryServiceInterface
interface
{
RegisterNode
(
ctx
context
.
Context
,
req
*
querypb
.
RegisterNodeRequest
)
(
*
querypb
.
RegisterNodeResponse
,
error
)
}
type
DataServiceInterface
interface
{
GetInsertBinlogPaths
(
ctx
context
.
Context
,
req
*
datapb
.
InsertBinlogPathRequest
)
(
*
datapb
.
InsertBinlogPathsResponse
,
error
)
GetSegmentStates
(
ctx
context
.
Context
,
req
*
datapb
.
SegmentStatesRequest
)
(
*
datapb
.
SegmentStatesResponse
,
error
)
}
type
IndexServiceInterface
interface
{
GetIndexFilePaths
(
ctx
context
.
Context
,
req
*
indexpb
.
IndexFilePathsRequest
)
(
*
indexpb
.
IndexFilePathsResponse
,
error
)
}
internal/types/types.go
浏览文件 @
7a3223dc
...
...
@@ -176,6 +176,7 @@ type QueryService interface {
Component
TimeTickProvider
RegisterNode
(
ctx
context
.
Context
,
req
*
querypb
.
RegisterNodeRequest
)
(
*
querypb
.
RegisterNodeResponse
,
error
)
ShowCollections
(
ctx
context
.
Context
,
req
*
querypb
.
ShowCollectionRequest
)
(
*
querypb
.
ShowCollectionResponse
,
error
)
LoadCollection
(
ctx
context
.
Context
,
req
*
querypb
.
LoadCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
ReleaseCollection
(
ctx
context
.
Context
,
req
*
querypb
.
ReleaseCollectionRequest
)
(
*
commonpb
.
Status
,
error
)
...
...
internal/util/typeutil/interface.go
浏览文件 @
7a3223dc
...
...
@@ -4,7 +4,6 @@ import (
"context"
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
"github.com/zilliztech/milvus-distributed/internal/proto/querypb"
...
...
@@ -24,24 +23,6 @@ type Component interface {
GetStatisticsChannel
(
ctx
context
.
Context
)
(
*
milvuspb
.
StringResponse
,
error
)
}
// TODO
type
IndexNodeInterface
interface
{
Service
Component
BuildIndex
(
ctx
context
.
Context
,
req
*
indexpb
.
BuildIndexCmd
)
(
*
commonpb
.
Status
,
error
)
DropIndex
(
ctx
context
.
Context
,
req
*
indexpb
.
DropIndexRequest
)
(
*
commonpb
.
Status
,
error
)
}
type
IndexServiceInterface
interface
{
Service
Component
RegisterNode
(
ctx
context
.
Context
,
req
*
indexpb
.
RegisterNodeRequest
)
(
*
indexpb
.
RegisterNodeResponse
,
error
)
BuildIndex
(
ctx
context
.
Context
,
req
*
indexpb
.
BuildIndexRequest
)
(
*
indexpb
.
BuildIndexResponse
,
error
)
GetIndexStates
(
ctx
context
.
Context
,
req
*
indexpb
.
IndexStatesRequest
)
(
*
indexpb
.
IndexStatesResponse
,
error
)
GetIndexFilePaths
(
ctx
context
.
Context
,
req
*
indexpb
.
IndexFilePathsRequest
)
(
*
indexpb
.
IndexFilePathsResponse
,
error
)
NotifyBuildIndex
(
ctx
context
.
Context
,
nty
*
indexpb
.
BuildIndexNotification
)
(
*
commonpb
.
Status
,
error
)
}
type
QueryServiceInterface
interface
{
Service
Component
...
...
tests/python_test/collection/test_collection_count.py
浏览文件 @
7a3223dc
此差异已折叠。
点击以展开。
tests/python_test/collection/test_collection_stats.py
浏览文件 @
7a3223dc
...
...
@@ -58,6 +58,7 @@ class TestGetCollectionStats:
def
insert_count
(
self
,
request
):
yield
request
.
param
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_name_not_existed
(
self
,
connect
,
collection
):
'''
target: get collection stats where collection name does not exist
...
...
@@ -69,6 +70,7 @@ class TestGetCollectionStats:
connect
.
get_collection_stats
(
collection_name
)
@
pytest
.
mark
.
level
(
2
)
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_name_invalid
(
self
,
connect
,
get_invalid_collection_name
):
'''
target: get collection stats where collection name is invalid
...
...
@@ -79,6 +81,7 @@ class TestGetCollectionStats:
with
pytest
.
raises
(
Exception
)
as
e
:
connect
.
get_collection_stats
(
collection_name
)
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_empty
(
self
,
connect
,
collection
):
'''
target: get collection stats where no entity in collection
...
...
@@ -89,6 +92,7 @@ class TestGetCollectionStats:
connect
.
flush
([
collection
])
assert
stats
[
row_count
]
==
0
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_without_connection
(
self
,
collection
,
dis_connect
):
'''
target: test count_entities, without connection
...
...
@@ -98,6 +102,7 @@ class TestGetCollectionStats:
with
pytest
.
raises
(
Exception
)
as
e
:
dis_connect
.
get_collection_stats
(
collection
)
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_batch
(
self
,
connect
,
collection
):
'''
target: get row count with collection_stats
...
...
@@ -110,6 +115,8 @@ class TestGetCollectionStats:
stats
=
connect
.
get_collection_stats
(
collection
)
assert
int
(
stats
[
row_count
])
==
default_nb
# @pytest.mark.tags("0331")
# TODO ci failed
def
test_get_collection_stats_single
(
self
,
connect
,
collection
):
'''
target: get row count with collection_stats
...
...
@@ -184,6 +191,7 @@ class TestGetCollectionStats:
# pdb.set_trace()
assert
compact_before
==
compact_after
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_partition
(
self
,
connect
,
collection
):
'''
target: get partition info in a collection
...
...
@@ -197,10 +205,11 @@ class TestGetCollectionStats:
stats
=
connect
.
get_collection_stats
(
collection
)
assert
stats
[
row_count
]
==
default_nb
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_partitions
(
self
,
connect
,
collection
):
'''
target: get partition info in a collection
method: create two partitions, add vectors in one of the partitions, call collection_stats and check
method: create two partitions, add vectors in one of the partitions, call collection_stats and check
expected: status ok, vectors added to one partition but not the other
'''
new_tag
=
"new_tag"
...
...
@@ -219,7 +228,7 @@ class TestGetCollectionStats:
stats
=
connect
.
get_collection_stats
(
collection
)
assert
stats
[
row_count
]
==
default_nb
*
3
#
@pytest.mark.tags("0331")
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_partitions_A
(
self
,
connect
,
collection
,
insert_count
):
'''
target: test collection rows_count is correct or not
...
...
@@ -236,7 +245,7 @@ class TestGetCollectionStats:
stats
=
connect
.
get_collection_stats
(
collection
)
assert
stats
[
row_count
]
==
insert_count
#
@pytest.mark.tags("0331")
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_partitions_B
(
self
,
connect
,
collection
,
insert_count
):
'''
target: test collection rows_count is correct or not
...
...
@@ -253,7 +262,7 @@ class TestGetCollectionStats:
stats
=
connect
.
get_collection_stats
(
collection
)
assert
stats
[
row_count
]
==
insert_count
#
@pytest.mark.tags("0331")
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_partitions_C
(
self
,
connect
,
collection
,
insert_count
):
'''
target: test collection rows_count is correct or not
...
...
@@ -271,7 +280,7 @@ class TestGetCollectionStats:
stats
=
connect
.
get_collection_stats
(
collection
)
assert
stats
[
row_count
]
==
insert_count
*
2
#
@pytest.mark.tags("0331")
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_partitions_D
(
self
,
connect
,
collection
,
insert_count
):
'''
target: test collection rows_count is correct or not
...
...
@@ -290,10 +299,11 @@ class TestGetCollectionStats:
assert
stats
[
row_count
]
==
insert_count
*
2
# TODO: assert metric type in stats response
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_after_index_created
(
self
,
connect
,
collection
,
get_simple_index
):
'''
target: test collection info after index created
method: create collection, add vectors, create index and call collection_stats
method: create collection, add vectors, create index and call collection_stats
expected: status ok, index created and shown in segments
'''
connect
.
insert
(
collection
,
default_entities
)
...
...
@@ -303,10 +313,11 @@ class TestGetCollectionStats:
assert
stats
[
row_count
]
==
default_nb
# TODO: assert metric type in stats response
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_after_index_created_ip
(
self
,
connect
,
collection
,
get_simple_index
):
'''
target: test collection info after index created
method: create collection, add vectors, create index and call collection_stats
method: create collection, add vectors, create index and call collection_stats
expected: status ok, index created and shown in segments
'''
get_simple_index
[
"metric_type"
]
=
"IP"
...
...
@@ -319,10 +330,11 @@ class TestGetCollectionStats:
assert
stats
[
row_count
]
==
default_nb
# TODO: assert metric type in stats response
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_after_index_created_jac
(
self
,
connect
,
binary_collection
,
get_jaccard_index
):
'''
target: test collection info after index created
method: create collection, add binary entities, create index and call collection_stats
method: create collection, add binary entities, create index and call collection_stats
expected: status ok, index created and shown in segments
'''
ids
=
connect
.
insert
(
binary_collection
,
default_binary_entities
)
...
...
@@ -331,10 +343,11 @@ class TestGetCollectionStats:
stats
=
connect
.
get_collection_stats
(
binary_collection
)
assert
stats
[
row_count
]
==
default_nb
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_get_collection_stats_after_create_different_index
(
self
,
connect
,
collection
):
'''
target: test collection info after index created repeatedly
method: create collection, add vectors, create index and call collection_stats multiple times
method: create collection, add vectors, create index and call collection_stats multiple times
expected: status ok, index info shown in segments
'''
ids
=
connect
.
insert
(
collection
,
default_entities
)
...
...
@@ -345,6 +358,7 @@ class TestGetCollectionStats:
stats
=
connect
.
get_collection_stats
(
collection
)
assert
stats
[
row_count
]
==
default_nb
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_collection_count_multi_collections
(
self
,
connect
):
'''
target: test collection rows_count is correct or not with multiple collections of L2
...
...
@@ -366,6 +380,7 @@ class TestGetCollectionStats:
connect
.
drop_collection
(
collection_list
[
i
])
@
pytest
.
mark
.
level
(
2
)
@
pytest
.
mark
.
tags
(
"0331"
)
def
test_collection_count_multi_collections_indexed
(
self
,
connect
):
'''
target: test collection rows_count is correct or not with multiple collections of L2
...
...
tests/python_test/utils.py
浏览文件 @
7a3223dc
...
...
@@ -871,7 +871,7 @@ def get_search_param(index_type, metric_type="L2"):
search_params
=
{
"metric_type"
:
metric_type
}
if
index_type
in
ivf
()
or
index_type
in
binary_support
():
search_params
.
update
({
"nprobe"
:
64
})
elif
index_type
in
[
"HNSW"
,
"RHNSW_SQ"
,
"RHNSW_PQ"
]:
elif
index_type
in
[
"HNSW"
,
"RHNSW_
FLAT"
,
"RHNSW_
SQ"
,
"RHNSW_PQ"
]:
search_params
.
update
({
"ef"
:
64
})
elif
index_type
==
"NSG"
:
search_params
.
update
({
"search_length"
:
100
})
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录