2021-01-15 14:38:36 +08:00
|
|
|
package indexnode
|
2020-12-10 17:55:55 +08:00
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
2021-01-12 18:03:24 +08:00
|
|
|
"fmt"
|
2020-12-13 06:48:05 +08:00
|
|
|
"log"
|
2020-12-22 08:14:36 +08:00
|
|
|
"strconv"
|
2020-12-13 06:48:05 +08:00
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
|
|
|
|
"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
|
|
|
|
|
2020-12-10 17:55:55 +08:00
|
|
|
"github.com/zilliztech/milvus-distributed/internal/errors"
|
2020-12-22 08:14:36 +08:00
|
|
|
"github.com/zilliztech/milvus-distributed/internal/kv"
|
2021-01-18 19:32:08 +08:00
|
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
|
2020-12-22 08:14:36 +08:00
|
|
|
"github.com/zilliztech/milvus-distributed/internal/storage"
|
2020-12-10 17:55:55 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
type task interface {
|
|
|
|
ID() UniqueID // return ReqID
|
|
|
|
SetID(uid UniqueID) // set ReqID
|
|
|
|
PreExecute() error
|
|
|
|
Execute() error
|
|
|
|
PostExecute() error
|
|
|
|
WaitToFinish() error
|
|
|
|
Notify(err error)
|
2020-12-13 06:48:05 +08:00
|
|
|
OnEnqueue() error
|
2020-12-10 17:55:55 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
type BaseTask struct {
|
2021-01-26 09:38:40 +08:00
|
|
|
done chan error
|
|
|
|
ctx context.Context
|
|
|
|
id UniqueID
|
2020-12-10 17:55:55 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (bt *BaseTask) ID() UniqueID {
|
|
|
|
return bt.id
|
|
|
|
}
|
|
|
|
|
|
|
|
func (bt *BaseTask) setID(id UniqueID) {
|
|
|
|
bt.id = id
|
|
|
|
}
|
|
|
|
|
|
|
|
func (bt *BaseTask) WaitToFinish() error {
|
2020-12-13 06:48:05 +08:00
|
|
|
select {
|
|
|
|
case <-bt.ctx.Done():
|
|
|
|
return errors.New("timeout")
|
|
|
|
case err := <-bt.done:
|
|
|
|
return err
|
2020-12-10 17:55:55 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (bt *BaseTask) Notify(err error) {
|
|
|
|
bt.done <- err
|
|
|
|
}
|
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
type IndexBuildTask struct {
|
2020-12-10 17:55:55 +08:00
|
|
|
BaseTask
|
2021-01-26 09:38:40 +08:00
|
|
|
index Index
|
|
|
|
kv kv.Base
|
|
|
|
savePaths []string
|
|
|
|
cmd *indexpb.BuildIndexCmd
|
|
|
|
serviceClient typeutil.IndexServiceInterface
|
|
|
|
nodeID UniqueID
|
|
|
|
}
|
|
|
|
|
|
|
|
func newIndexBuildTask() *IndexBuildTask {
|
|
|
|
ctx := context.Background()
|
|
|
|
return &IndexBuildTask{
|
|
|
|
BaseTask: BaseTask{
|
|
|
|
ctx: ctx,
|
|
|
|
done: make(chan error), // intend to do this
|
|
|
|
},
|
|
|
|
}
|
2020-12-10 17:55:55 +08:00
|
|
|
}
|
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
func (it *IndexBuildTask) SetID(ID UniqueID) {
|
2020-12-13 06:48:05 +08:00
|
|
|
it.BaseTask.setID(ID)
|
|
|
|
}
|
2020-12-10 17:55:55 +08:00
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
func (it *IndexBuildTask) OnEnqueue() error {
|
2021-02-02 19:56:04 +08:00
|
|
|
it.SetID(it.cmd.IndexBuildID)
|
2021-01-26 09:38:40 +08:00
|
|
|
log.Printf("[IndexBuilderTask] Enqueue TaskID: %v", it.ID())
|
2020-12-10 17:55:55 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
func (it *IndexBuildTask) PreExecute() error {
|
|
|
|
log.Println("preExecute...")
|
2020-12-10 17:55:55 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
func (it *IndexBuildTask) PostExecute() error {
|
|
|
|
log.Println("PostExecute...")
|
|
|
|
var err error
|
|
|
|
defer func() {
|
|
|
|
if err != nil {
|
|
|
|
it.Rollback()
|
2020-12-13 06:48:05 +08:00
|
|
|
}
|
2021-01-26 09:38:40 +08:00
|
|
|
}()
|
2020-12-13 06:48:05 +08:00
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
if it.serviceClient == nil {
|
|
|
|
err = errors.New("IndexBuildTask, serviceClient is nil")
|
|
|
|
return err
|
|
|
|
}
|
2020-12-10 17:55:55 +08:00
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
nty := &indexpb.BuildIndexNotification{
|
|
|
|
Status: &commonpb.Status{
|
|
|
|
ErrorCode: commonpb.ErrorCode_SUCCESS,
|
2020-12-13 06:48:05 +08:00
|
|
|
},
|
2021-02-02 19:56:04 +08:00
|
|
|
IndexBuildID: it.cmd.IndexBuildID,
|
2021-01-26 09:38:40 +08:00
|
|
|
NodeID: it.nodeID,
|
|
|
|
IndexFilePaths: it.savePaths,
|
2020-12-13 06:48:05 +08:00
|
|
|
}
|
2020-12-10 17:55:55 +08:00
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
resp, err := it.serviceClient.NotifyBuildIndex(nty)
|
|
|
|
if err != nil {
|
|
|
|
log.Println("IndexBuildTask notify err:", err.Error())
|
|
|
|
return err
|
2020-12-13 06:48:05 +08:00
|
|
|
}
|
2020-12-10 17:55:55 +08:00
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
if resp.ErrorCode != commonpb.ErrorCode_SUCCESS {
|
|
|
|
err = errors.New(resp.Reason)
|
|
|
|
}
|
|
|
|
return err
|
2020-12-13 06:48:05 +08:00
|
|
|
}
|
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
func (it *IndexBuildTask) Rollback() error {
|
2020-12-10 17:55:55 +08:00
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
if it.savePaths == nil {
|
|
|
|
return nil
|
|
|
|
}
|
2020-12-13 06:48:05 +08:00
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
err := it.kv.MultiRemove(it.savePaths)
|
2020-12-13 06:48:05 +08:00
|
|
|
if err != nil {
|
2021-01-26 09:38:40 +08:00
|
|
|
log.Println("IndexBuildTask Rollback Failed:", err.Error())
|
2020-12-13 06:48:05 +08:00
|
|
|
return err
|
|
|
|
}
|
2021-01-26 09:38:40 +08:00
|
|
|
return nil
|
|
|
|
}
|
2020-12-22 08:14:36 +08:00
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
func (it *IndexBuildTask) Execute() error {
|
|
|
|
log.Println("start build index ...")
|
|
|
|
var err error
|
2020-12-22 08:14:36 +08:00
|
|
|
typeParams := make(map[string]string)
|
2021-01-26 09:38:40 +08:00
|
|
|
for _, kvPair := range it.cmd.Req.GetTypeParams() {
|
2020-12-22 08:14:36 +08:00
|
|
|
key, value := kvPair.GetKey(), kvPair.GetValue()
|
|
|
|
_, ok := typeParams[key]
|
|
|
|
if ok {
|
|
|
|
return errors.New("duplicated key in type params")
|
|
|
|
}
|
|
|
|
typeParams[key] = value
|
|
|
|
}
|
|
|
|
|
|
|
|
indexParams := make(map[string]string)
|
2021-01-26 09:38:40 +08:00
|
|
|
for _, kvPair := range it.cmd.Req.GetIndexParams() {
|
2020-12-22 08:14:36 +08:00
|
|
|
key, value := kvPair.GetKey(), kvPair.GetValue()
|
|
|
|
_, ok := indexParams[key]
|
|
|
|
if ok {
|
|
|
|
return errors.New("duplicated key in index params")
|
|
|
|
}
|
|
|
|
indexParams[key] = value
|
|
|
|
}
|
|
|
|
|
|
|
|
it.index, err = NewCIndex(typeParams, indexParams)
|
|
|
|
if err != nil {
|
2021-01-23 20:58:46 +08:00
|
|
|
fmt.Println("NewCIndex err:", err.Error())
|
2020-12-22 08:14:36 +08:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
getKeyByPathNaive := func(path string) string {
|
|
|
|
// splitElements := strings.Split(path, "/")
|
|
|
|
// return splitElements[len(splitElements)-1]
|
|
|
|
return path
|
|
|
|
}
|
|
|
|
getValueByPath := func(path string) ([]byte, error) {
|
|
|
|
data, err := it.kv.Load(path)
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
return []byte(data), nil
|
|
|
|
}
|
|
|
|
getBlobByPath := func(path string) (*Blob, error) {
|
|
|
|
value, err := getValueByPath(path)
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
return &Blob{
|
|
|
|
Key: getKeyByPathNaive(path),
|
|
|
|
Value: value,
|
|
|
|
}, nil
|
|
|
|
}
|
|
|
|
getStorageBlobs := func(blobs []*Blob) []*storage.Blob {
|
2020-12-25 11:10:31 +08:00
|
|
|
return blobs
|
2020-12-22 08:14:36 +08:00
|
|
|
}
|
|
|
|
|
2021-01-26 09:38:40 +08:00
|
|
|
toLoadDataPaths := it.cmd.Req.GetDataPaths()
|
2020-12-22 08:14:36 +08:00
|
|
|
keys := make([]string, 0)
|
|
|
|
blobs := make([]*Blob, 0)
|
|
|
|
for _, path := range toLoadDataPaths {
|
|
|
|
keys = append(keys, getKeyByPathNaive(path))
|
|
|
|
blob, err := getBlobByPath(path)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
blobs = append(blobs, blob)
|
|
|
|
}
|
|
|
|
|
|
|
|
storageBlobs := getStorageBlobs(blobs)
|
|
|
|
var insertCodec storage.InsertCodec
|
2021-01-15 14:38:36 +08:00
|
|
|
partitionID, segmentID, insertData, err2 := insertCodec.Deserialize(storageBlobs)
|
2021-01-23 20:58:46 +08:00
|
|
|
//fmt.Println("IndexBuilder for segmentID,", segmentID)
|
2021-01-15 14:38:36 +08:00
|
|
|
if err2 != nil {
|
|
|
|
return err2
|
|
|
|
}
|
2020-12-22 08:14:36 +08:00
|
|
|
if len(insertData.Data) != 1 {
|
|
|
|
return errors.New("we expect only one field in deserialized insert data")
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, value := range insertData.Data {
|
|
|
|
// TODO: BinaryVectorFieldData
|
2021-01-07 15:39:20 +08:00
|
|
|
floatVectorFieldData, fOk := value.(*storage.FloatVectorFieldData)
|
|
|
|
if fOk {
|
|
|
|
err = it.index.BuildFloatVecIndexWithoutIds(floatVectorFieldData.Data)
|
|
|
|
if err != nil {
|
2021-01-23 20:58:46 +08:00
|
|
|
fmt.Println("BuildFloatVecIndexWithoutIds, error:", err.Error())
|
2021-01-07 15:39:20 +08:00
|
|
|
return err
|
|
|
|
}
|
2021-01-07 14:56:17 +08:00
|
|
|
}
|
|
|
|
|
2021-01-07 15:39:20 +08:00
|
|
|
binaryVectorFieldData, bOk := value.(*storage.BinaryVectorFieldData)
|
|
|
|
if bOk {
|
|
|
|
err = it.index.BuildBinaryVecIndexWithoutIds(binaryVectorFieldData.Data)
|
|
|
|
if err != nil {
|
2021-01-23 20:58:46 +08:00
|
|
|
fmt.Println("BuildBinaryVecIndexWithoutIds, err:", err.Error())
|
2021-01-07 15:39:20 +08:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-01-07 16:13:28 +08:00
|
|
|
if !fOk && !bOk {
|
2021-01-07 15:39:20 +08:00
|
|
|
return errors.New("we expect FloatVectorFieldData or BinaryVectorFieldData")
|
2020-12-22 08:14:36 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
indexBlobs, err := it.index.Serialize()
|
|
|
|
if err != nil {
|
2021-01-23 20:58:46 +08:00
|
|
|
fmt.Println("serialize ... err:", err.Error())
|
|
|
|
|
2020-12-22 08:14:36 +08:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
var indexCodec storage.IndexCodec
|
2021-02-02 19:56:11 +08:00
|
|
|
serializedIndexBlobs, err := indexCodec.Serialize(getStorageBlobs(indexBlobs), indexParams, it.cmd.Req.IndexName, it.cmd.Req.IndexID)
|
2020-12-22 08:14:36 +08:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
getSavePathByKey := func(key string) string {
|
|
|
|
// TODO: fix me, use more reasonable method
|
2021-02-02 19:56:04 +08:00
|
|
|
return strconv.Itoa(int(it.cmd.IndexBuildID)) + "/" + strconv.Itoa(int(partitionID)) + "/" + strconv.Itoa(int(segmentID)) + "/" + key
|
2020-12-22 08:14:36 +08:00
|
|
|
}
|
|
|
|
saveBlob := func(path string, value []byte) error {
|
|
|
|
return it.kv.Save(path, string(value))
|
|
|
|
}
|
|
|
|
|
|
|
|
it.savePaths = make([]string, 0)
|
|
|
|
for _, blob := range serializedIndexBlobs {
|
2020-12-25 11:10:31 +08:00
|
|
|
key, value := blob.Key, blob.Value
|
2020-12-22 08:14:36 +08:00
|
|
|
savePath := getSavePathByKey(key)
|
|
|
|
err := saveBlob(savePath, value)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
it.savePaths = append(it.savePaths, savePath)
|
|
|
|
}
|
|
|
|
}
|
2021-01-26 09:38:40 +08:00
|
|
|
err = it.index.Delete()
|
|
|
|
if err != nil {
|
|
|
|
log.Print("CIndexDelete Failed")
|
|
|
|
}
|
2021-01-23 20:58:46 +08:00
|
|
|
return nil
|
2020-12-10 17:55:55 +08:00
|
|
|
}
|