mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 05:18:52 +08:00
a7b27db63e
* update fillRetrieveResults parameter Signed-off-by: yudong.cai <yudong.cai@zilliz.com> * optimize fillVectorFieldData process Signed-off-by: yudong.cai <yudong.cai@zilliz.com>
367 lines
10 KiB
Go
367 lines
10 KiB
Go
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software distributed under the License
|
|
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
|
|
// or implied. See the License for the specific language governing permissions and limitations under the License.
|
|
|
|
package querynode
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
|
|
"github.com/golang/protobuf/proto"
|
|
"go.uber.org/zap"
|
|
|
|
"github.com/milvus-io/milvus/internal/kv"
|
|
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
|
|
minioKV "github.com/milvus-io/milvus/internal/kv/minio"
|
|
"github.com/milvus-io/milvus/internal/log"
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
|
"github.com/milvus-io/milvus/internal/proto/querypb"
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
"github.com/milvus-io/milvus/internal/types"
|
|
"github.com/milvus-io/milvus/internal/util/funcutil"
|
|
)
|
|
|
|
const (
|
|
queryCoordSegmentMetaPrefix = "queryCoord-segmentMeta"
|
|
queryNodeSegmentMetaPrefix = "queryNode-segmentMeta"
|
|
)
|
|
|
|
// segmentLoader is only responsible for loading the field data from binlog
|
|
type segmentLoader struct {
|
|
historicalReplica ReplicaInterface
|
|
|
|
dataCoord types.DataCoord
|
|
|
|
minioKV kv.BaseKV // minio minioKV
|
|
etcdKV *etcdkv.EtcdKV
|
|
|
|
indexLoader *indexLoader
|
|
}
|
|
|
|
func (loader *segmentLoader) loadSegmentOfConditionHandOff(req *querypb.LoadSegmentsRequest) error {
|
|
return errors.New("TODO: implement hand off")
|
|
}
|
|
|
|
func (loader *segmentLoader) loadSegmentOfConditionLoadBalance(req *querypb.LoadSegmentsRequest) error {
|
|
return loader.loadSegment(req, false)
|
|
}
|
|
|
|
func (loader *segmentLoader) loadSegmentOfConditionGRPC(req *querypb.LoadSegmentsRequest) error {
|
|
return loader.loadSegment(req, true)
|
|
}
|
|
|
|
func (loader *segmentLoader) loadSegmentOfConditionNodeDown(req *querypb.LoadSegmentsRequest) error {
|
|
return loader.loadSegment(req, true)
|
|
}
|
|
|
|
func (loader *segmentLoader) loadSegment(req *querypb.LoadSegmentsRequest, onService bool) error {
|
|
// no segment needs to load, return
|
|
if len(req.Infos) == 0 {
|
|
return nil
|
|
}
|
|
|
|
// start to load
|
|
for _, info := range req.Infos {
|
|
segmentID := info.SegmentID
|
|
partitionID := info.PartitionID
|
|
collectionID := info.CollectionID
|
|
|
|
// init replica
|
|
hasCollectionInHistorical := loader.historicalReplica.hasCollection(collectionID)
|
|
hasPartitionInHistorical := loader.historicalReplica.hasPartition(partitionID)
|
|
if !hasCollectionInHistorical {
|
|
err := loader.historicalReplica.addCollection(collectionID, req.Schema)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if !hasPartitionInHistorical {
|
|
err := loader.historicalReplica.addPartition(collectionID, partitionID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
collection, err := loader.historicalReplica.getCollectionByID(collectionID)
|
|
if err != nil {
|
|
log.Warn(err.Error())
|
|
continue
|
|
}
|
|
segment := newSegment(collection, segmentID, partitionID, collectionID, "", segmentTypeSealed, onService)
|
|
err = loader.loadSegmentInternal(collectionID, segment, info)
|
|
if err != nil {
|
|
deleteSegment(segment)
|
|
log.Error(err.Error())
|
|
continue
|
|
}
|
|
err = loader.historicalReplica.setSegment(segment)
|
|
if err != nil {
|
|
deleteSegment(segment)
|
|
log.Error(err.Error())
|
|
continue
|
|
}
|
|
if onService {
|
|
key := fmt.Sprintf("%s/%d", queryCoordSegmentMetaPrefix, segmentID)
|
|
value, err := loader.etcdKV.Load(key)
|
|
if err != nil {
|
|
deleteSegment(segment)
|
|
log.Error("error when load segment info from etcd", zap.Any("error", err.Error()))
|
|
continue
|
|
}
|
|
segmentInfo := &querypb.SegmentInfo{}
|
|
err = proto.UnmarshalText(value, segmentInfo)
|
|
if err != nil {
|
|
deleteSegment(segment)
|
|
log.Error("error when unmarshal segment info from etcd", zap.Any("error", err.Error()))
|
|
continue
|
|
}
|
|
segmentInfo.SegmentState = querypb.SegmentState_sealed
|
|
newKey := fmt.Sprintf("%s/%d", queryNodeSegmentMetaPrefix, segmentID)
|
|
err = loader.etcdKV.Save(newKey, proto.MarshalTextString(segmentInfo))
|
|
if err != nil {
|
|
deleteSegment(segment)
|
|
log.Error("error when update segment info to etcd", zap.Any("error", err.Error()))
|
|
}
|
|
}
|
|
}
|
|
|
|
// sendQueryNodeStats
|
|
return loader.indexLoader.sendQueryNodeStats()
|
|
}
|
|
|
|
func (loader *segmentLoader) loadSegmentInternal(collectionID UniqueID, segment *Segment, segmentLoadInfo *querypb.SegmentLoadInfo) error {
|
|
vectorFieldIDs, err := loader.historicalReplica.getVecFieldIDsByCollectionID(collectionID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(vectorFieldIDs) <= 0 {
|
|
return fmt.Errorf("no vector field in collection %d", collectionID)
|
|
}
|
|
|
|
// add VectorFieldInfo for vector fields
|
|
for _, fieldBinlog := range segmentLoadInfo.BinlogPaths {
|
|
if funcutil.SliceContain(vectorFieldIDs, fieldBinlog.FieldID) {
|
|
vectorFieldInfo := newVectorFieldInfo(fieldBinlog)
|
|
segment.setVectorFieldInfo(fieldBinlog.FieldID, vectorFieldInfo)
|
|
}
|
|
}
|
|
|
|
indexedFieldIDs := make([]int64, 0)
|
|
for _, vecFieldID := range vectorFieldIDs {
|
|
err = loader.indexLoader.setIndexInfo(collectionID, segment, vecFieldID)
|
|
if err != nil {
|
|
log.Warn(err.Error())
|
|
continue
|
|
}
|
|
indexedFieldIDs = append(indexedFieldIDs, vecFieldID)
|
|
}
|
|
|
|
// we don't need to load raw data for indexed vector field
|
|
fieldBinlogs := loader.filterFieldBinlogs(segmentLoadInfo.BinlogPaths, indexedFieldIDs)
|
|
|
|
log.Debug("loading insert...")
|
|
err = loader.loadSegmentFieldsData(segment, fieldBinlogs)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, id := range indexedFieldIDs {
|
|
log.Debug("loading index...")
|
|
err = loader.indexLoader.loadIndex(segment, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
//func (loader *segmentLoader) GetSegmentStates(segmentID UniqueID) (*datapb.GetSegmentStatesResponse, error) {
|
|
// ctx := context.TODO()
|
|
// if loader.dataCoord == nil {
|
|
// return nil, errors.New("null data service client")
|
|
// }
|
|
//
|
|
// segmentStatesRequest := &datapb.GetSegmentStatesRequest{
|
|
// SegmentIDs: []int64{segmentID},
|
|
// }
|
|
// statesResponse, err := loader.dataCoord.GetSegmentStates(ctx, segmentStatesRequest)
|
|
// if err != nil || statesResponse.Status.ErrorCode != commonpb.ErrorCode_Success {
|
|
// return nil, err
|
|
// }
|
|
// if len(statesResponse.States) != 1 {
|
|
// return nil, errors.New("segment states' len should be 1")
|
|
// }
|
|
//
|
|
// return statesResponse, nil
|
|
//}
|
|
|
|
func (loader *segmentLoader) filterFieldBinlogs(fieldBinlogs []*datapb.FieldBinlog, skipFieldIDs []int64) []*datapb.FieldBinlog {
|
|
result := make([]*datapb.FieldBinlog, 0)
|
|
for _, fieldBinlog := range fieldBinlogs {
|
|
if !funcutil.SliceContain(skipFieldIDs, fieldBinlog.FieldID) {
|
|
result = append(result, fieldBinlog)
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
func (loader *segmentLoader) loadSegmentFieldsData(segment *Segment, fieldBinlogs []*datapb.FieldBinlog) error {
|
|
iCodec := storage.InsertCodec{}
|
|
defer func() {
|
|
err := iCodec.Close()
|
|
if err != nil {
|
|
log.Error(err.Error())
|
|
}
|
|
}()
|
|
blobs := make([]*storage.Blob, 0)
|
|
for _, fb := range fieldBinlogs {
|
|
log.Debug("load segment fields data",
|
|
zap.Int64("segmentID", segment.segmentID),
|
|
zap.Any("fieldID", fb.FieldID),
|
|
zap.String("paths", fmt.Sprintln(fb.Binlogs)),
|
|
)
|
|
for _, path := range fb.Binlogs {
|
|
p := path
|
|
binLog, err := loader.minioKV.Load(path)
|
|
if err != nil {
|
|
// TODO: return or continue?
|
|
return err
|
|
}
|
|
blob := &storage.Blob{
|
|
Key: p,
|
|
Value: []byte(binLog),
|
|
}
|
|
blobs = append(blobs, blob)
|
|
}
|
|
// mark the flag that vector raw data will be loaded into memory
|
|
if vecFieldInfo, err := segment.getVectorFieldInfo(fb.FieldID); err == nil {
|
|
vecFieldInfo.setRawDataInMemory(true)
|
|
}
|
|
}
|
|
|
|
_, _, insertData, err := iCodec.Deserialize(blobs)
|
|
if err != nil {
|
|
log.Error(err.Error())
|
|
return err
|
|
}
|
|
|
|
for fieldID, value := range insertData.Data {
|
|
var numRows int
|
|
var data interface{}
|
|
switch fieldData := value.(type) {
|
|
case *storage.BoolFieldData:
|
|
numRows = fieldData.NumRows
|
|
data = fieldData.Data
|
|
case *storage.Int8FieldData:
|
|
numRows = fieldData.NumRows
|
|
data = fieldData.Data
|
|
case *storage.Int16FieldData:
|
|
numRows = fieldData.NumRows
|
|
data = fieldData.Data
|
|
case *storage.Int32FieldData:
|
|
numRows = fieldData.NumRows
|
|
data = fieldData.Data
|
|
case *storage.Int64FieldData:
|
|
numRows = fieldData.NumRows
|
|
data = fieldData.Data
|
|
case *storage.FloatFieldData:
|
|
numRows = fieldData.NumRows
|
|
data = fieldData.Data
|
|
case *storage.DoubleFieldData:
|
|
numRows = fieldData.NumRows
|
|
data = fieldData.Data
|
|
case storage.StringFieldData:
|
|
numRows = fieldData.NumRows
|
|
data = fieldData.Data
|
|
case *storage.FloatVectorFieldData:
|
|
numRows = fieldData.NumRows
|
|
data = fieldData.Data
|
|
case *storage.BinaryVectorFieldData:
|
|
numRows = fieldData.NumRows
|
|
data = fieldData.Data
|
|
default:
|
|
return errors.New("unexpected field data type")
|
|
}
|
|
err = segment.segmentLoadFieldData(fieldID, numRows, data)
|
|
if err != nil {
|
|
// TODO: return or continue?
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (loader *segmentLoader) loadSegmentVectorFieldData(info *VectorFieldInfo) error {
|
|
iCodec := storage.InsertCodec{}
|
|
defer func() {
|
|
err := iCodec.Close()
|
|
if err != nil {
|
|
log.Error(err.Error())
|
|
}
|
|
}()
|
|
for _, path := range info.fieldBinlog.Binlogs {
|
|
if data := info.getRawData(path); data != nil {
|
|
continue
|
|
}
|
|
|
|
log.Debug("load vector raw data", zap.String("path", path))
|
|
|
|
binLog, err := loader.minioKV.Load(path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
blob := &storage.Blob{
|
|
Key: path,
|
|
Value: []byte(binLog),
|
|
}
|
|
|
|
insertFieldData, err := iCodec.DeserializeOneVectorBinlog(blob)
|
|
if err != nil {
|
|
log.Error(err.Error())
|
|
return err
|
|
}
|
|
|
|
// save raw data into segment.vectorFieldInfo
|
|
info.setRawData(path, insertFieldData.Data)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func newSegmentLoader(ctx context.Context, rootCoord types.RootCoord, indexCoord types.IndexCoord, replica ReplicaInterface, etcdKV *etcdkv.EtcdKV) *segmentLoader {
|
|
option := &minioKV.Option{
|
|
Address: Params.MinioEndPoint,
|
|
AccessKeyID: Params.MinioAccessKeyID,
|
|
SecretAccessKeyID: Params.MinioSecretAccessKey,
|
|
UseSSL: Params.MinioUseSSLStr,
|
|
CreateBucket: true,
|
|
BucketName: Params.MinioBucketName,
|
|
}
|
|
|
|
client, err := minioKV.NewMinIOKV(ctx, option)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
iLoader := newIndexLoader(ctx, rootCoord, indexCoord, replica)
|
|
return &segmentLoader{
|
|
historicalReplica: replica,
|
|
|
|
minioKV: client,
|
|
etcdKV: etcdKV,
|
|
|
|
indexLoader: iLoader,
|
|
}
|
|
}
|