mirror of https://github.com/milvus-io/milvus.git
319 lines
9.1 KiB
Go
319 lines
9.1 KiB
Go
// Licensed to the LF AI & Data foundation under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you 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 syncmgr
|
|
|
|
import (
|
|
"context"
|
|
"path"
|
|
|
|
"github.com/samber/lo"
|
|
"go.uber.org/zap"
|
|
|
|
"github.com/milvus-io/milvus/internal/allocator"
|
|
"github.com/milvus-io/milvus/internal/flushcommon/metacache"
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
"github.com/milvus-io/milvus/pkg/v2/common"
|
|
"github.com/milvus-io/milvus/pkg/v2/log"
|
|
"github.com/milvus-io/milvus/pkg/v2/proto/datapb"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/metautil"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/retry"
|
|
)
|
|
|
|
type PackWriter interface {
|
|
Write(ctx context.Context, pack *SyncPack) (
|
|
inserts []*datapb.Binlog, deletes *datapb.Binlog, stats *datapb.Binlog, bm25Stats *datapb.Binlog,
|
|
size int64, err error)
|
|
}
|
|
|
|
type BulkPackWriter struct {
|
|
metaCache metacache.MetaCache
|
|
chunkManager storage.ChunkManager
|
|
allocator allocator.Interface
|
|
writeRetryOpts []retry.Option
|
|
|
|
// prefetched log ids
|
|
ids []int64
|
|
sizeWritten int64
|
|
}
|
|
|
|
func NewBulkPackWriter(metaCache metacache.MetaCache, chunkManager storage.ChunkManager,
|
|
allocator allocator.Interface, writeRetryOpts ...retry.Option,
|
|
) *BulkPackWriter {
|
|
return &BulkPackWriter{
|
|
metaCache: metaCache,
|
|
chunkManager: chunkManager,
|
|
allocator: allocator,
|
|
writeRetryOpts: writeRetryOpts,
|
|
}
|
|
}
|
|
|
|
func (bw *BulkPackWriter) Write(ctx context.Context, pack *SyncPack) (
|
|
inserts map[int64]*datapb.FieldBinlog,
|
|
deltas *datapb.FieldBinlog,
|
|
stats map[int64]*datapb.FieldBinlog,
|
|
bm25Stats map[int64]*datapb.FieldBinlog,
|
|
size int64,
|
|
err error,
|
|
) {
|
|
err = bw.prefetchIDs(pack)
|
|
if err != nil {
|
|
log.Warn("failed allocate ids for sync task", zap.Error(err))
|
|
return
|
|
}
|
|
|
|
if inserts, err = bw.writeInserts(ctx, pack); err != nil {
|
|
log.Error("failed to write insert data", zap.Error(err))
|
|
return
|
|
}
|
|
if stats, err = bw.writeStats(ctx, pack); err != nil {
|
|
log.Error("failed to process stats blob", zap.Error(err))
|
|
return
|
|
}
|
|
if deltas, err = bw.writeDelta(ctx, pack); err != nil {
|
|
log.Error("failed to process delta blob", zap.Error(err))
|
|
return
|
|
}
|
|
if bm25Stats, err = bw.writeBM25Stasts(ctx, pack); err != nil {
|
|
log.Error("failed to process bm25 stats blob", zap.Error(err))
|
|
return
|
|
}
|
|
|
|
size = bw.sizeWritten
|
|
|
|
return
|
|
}
|
|
|
|
// prefetchIDs pre-allcates ids depending on the number of blobs current task contains.
|
|
func (bw *BulkPackWriter) prefetchIDs(pack *SyncPack) error {
|
|
totalIDCount := 0
|
|
if len(pack.insertData) > 0 {
|
|
totalIDCount += len(pack.insertData[0].Data) * 2 // binlogs and statslogs
|
|
}
|
|
if pack.isFlush {
|
|
totalIDCount++ // merged stats log
|
|
}
|
|
if pack.deltaData != nil {
|
|
totalIDCount++
|
|
}
|
|
if pack.bm25Stats != nil {
|
|
totalIDCount += len(pack.bm25Stats)
|
|
if pack.isFlush {
|
|
totalIDCount++ // merged bm25 stats
|
|
}
|
|
}
|
|
|
|
if totalIDCount == 0 {
|
|
return nil
|
|
}
|
|
start, _, err := bw.allocator.Alloc(uint32(totalIDCount))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
bw.ids = lo.RangeFrom(start, totalIDCount)
|
|
return nil
|
|
}
|
|
|
|
func (bw *BulkPackWriter) nextID() int64 {
|
|
if len(bw.ids) == 0 {
|
|
panic("pre-fetched ids exhausted")
|
|
}
|
|
r := bw.ids[0]
|
|
bw.ids = bw.ids[1:]
|
|
return r
|
|
}
|
|
|
|
func (bw *BulkPackWriter) writeLog(ctx context.Context, blob *storage.Blob,
|
|
root, p string, pack *SyncPack,
|
|
) (*datapb.Binlog, error) {
|
|
key := path.Join(bw.chunkManager.RootPath(), root, p)
|
|
err := retry.Do(ctx, func() error {
|
|
return bw.chunkManager.Write(ctx, key, blob.Value)
|
|
}, bw.writeRetryOpts...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
size := int64(len(blob.GetValue()))
|
|
bw.sizeWritten += size
|
|
return &datapb.Binlog{
|
|
EntriesNum: blob.RowNum,
|
|
TimestampFrom: pack.tsFrom,
|
|
TimestampTo: pack.tsTo,
|
|
LogPath: key,
|
|
LogSize: size,
|
|
MemorySize: blob.MemorySize,
|
|
}, nil
|
|
}
|
|
|
|
func (bw *BulkPackWriter) writeInserts(ctx context.Context, pack *SyncPack) (map[int64]*datapb.FieldBinlog, error) {
|
|
if len(pack.insertData) == 0 {
|
|
return make(map[int64]*datapb.FieldBinlog), nil
|
|
}
|
|
|
|
serializer, err := NewStorageSerializer(bw.metaCache)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
binlogBlobs, err := serializer.serializeBinlog(ctx, pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
logs := make(map[int64]*datapb.FieldBinlog)
|
|
for fieldID, blob := range binlogBlobs {
|
|
k := metautil.JoinIDPath(pack.collectionID, pack.partitionID, pack.segmentID, fieldID, bw.nextID())
|
|
binlog, err := bw.writeLog(ctx, blob, common.SegmentInsertLogPath, k, pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
logs[fieldID] = &datapb.FieldBinlog{
|
|
FieldID: fieldID,
|
|
Binlogs: []*datapb.Binlog{binlog},
|
|
}
|
|
}
|
|
return logs, nil
|
|
}
|
|
|
|
func (bw *BulkPackWriter) writeStats(ctx context.Context, pack *SyncPack) (map[int64]*datapb.FieldBinlog, error) {
|
|
if len(pack.insertData) == 0 {
|
|
return make(map[int64]*datapb.FieldBinlog), nil
|
|
}
|
|
serializer, err := NewStorageSerializer(bw.metaCache)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
singlePKStats, batchStatsBlob, err := serializer.serializeStatslog(pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
actions := []metacache.SegmentAction{metacache.RollStats(singlePKStats)}
|
|
bw.metaCache.UpdateSegments(metacache.MergeSegmentAction(actions...), metacache.WithSegmentIDs(pack.segmentID))
|
|
|
|
pkFieldID := serializer.pkField.GetFieldID()
|
|
binlogs := make([]*datapb.Binlog, 0)
|
|
k := metautil.JoinIDPath(pack.collectionID, pack.partitionID, pack.segmentID, pkFieldID, bw.nextID())
|
|
if binlog, err := bw.writeLog(ctx, batchStatsBlob, common.SegmentStatslogPath, k, pack); err != nil {
|
|
return nil, err
|
|
} else {
|
|
binlogs = append(binlogs, binlog)
|
|
}
|
|
|
|
if pack.isFlush && pack.level != datapb.SegmentLevel_L0 {
|
|
mergedStatsBlob, err := serializer.serializeMergedPkStats(pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
k := metautil.JoinIDPath(pack.collectionID, pack.partitionID, pack.segmentID, pkFieldID, int64(storage.CompoundStatsType))
|
|
binlog, err := bw.writeLog(ctx, mergedStatsBlob, common.SegmentStatslogPath, k, pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
binlogs = append(binlogs, binlog)
|
|
}
|
|
|
|
logs := make(map[int64]*datapb.FieldBinlog)
|
|
logs[pkFieldID] = &datapb.FieldBinlog{
|
|
FieldID: pkFieldID,
|
|
Binlogs: binlogs,
|
|
}
|
|
return logs, nil
|
|
}
|
|
|
|
func (bw *BulkPackWriter) writeBM25Stasts(ctx context.Context, pack *SyncPack) (map[int64]*datapb.FieldBinlog, error) {
|
|
if len(pack.bm25Stats) == 0 {
|
|
return make(map[int64]*datapb.FieldBinlog), nil
|
|
}
|
|
|
|
serializer, err := NewStorageSerializer(bw.metaCache)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
bm25Blobs, err := serializer.serializeBM25Stats(pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
logs := make(map[int64]*datapb.FieldBinlog)
|
|
for fieldID, blob := range bm25Blobs {
|
|
k := metautil.JoinIDPath(pack.collectionID, pack.partitionID, pack.segmentID, fieldID, bw.nextID())
|
|
binlog, err := bw.writeLog(ctx, blob, common.SegmentBm25LogPath, k, pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
logs[fieldID] = &datapb.FieldBinlog{
|
|
FieldID: fieldID,
|
|
Binlogs: []*datapb.Binlog{binlog},
|
|
}
|
|
}
|
|
|
|
actions := []metacache.SegmentAction{metacache.MergeBm25Stats(pack.bm25Stats)}
|
|
bw.metaCache.UpdateSegments(metacache.MergeSegmentAction(actions...), metacache.WithSegmentIDs(pack.segmentID))
|
|
|
|
if pack.isFlush {
|
|
if pack.level != datapb.SegmentLevel_L0 {
|
|
if hasBM25Function(bw.metaCache.Schema()) {
|
|
mergedBM25Blob, err := serializer.serializeMergedBM25Stats(pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for fieldID, blob := range mergedBM25Blob {
|
|
k := metautil.JoinIDPath(pack.collectionID, pack.partitionID, pack.segmentID, fieldID, int64(storage.CompoundStatsType))
|
|
binlog, err := bw.writeLog(ctx, blob, common.SegmentBm25LogPath, k, pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
fieldBinlog, ok := logs[fieldID]
|
|
if !ok {
|
|
fieldBinlog = &datapb.FieldBinlog{
|
|
FieldID: fieldID,
|
|
}
|
|
logs[fieldID] = fieldBinlog
|
|
}
|
|
fieldBinlog.Binlogs = append(fieldBinlog.Binlogs, binlog)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return logs, nil
|
|
}
|
|
|
|
func (bw *BulkPackWriter) writeDelta(ctx context.Context, pack *SyncPack) (*datapb.FieldBinlog, error) {
|
|
if pack.deltaData == nil {
|
|
return &datapb.FieldBinlog{}, nil
|
|
}
|
|
s, err := NewStorageSerializer(bw.metaCache)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
deltaBlob, err := s.serializeDeltalog(pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
k := metautil.JoinIDPath(pack.collectionID, pack.partitionID, pack.segmentID, bw.nextID())
|
|
deltalog, err := bw.writeLog(ctx, deltaBlob, common.SegmentDeltaLogPath, k, pack)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &datapb.FieldBinlog{
|
|
FieldID: s.pkField.GetFieldID(),
|
|
Binlogs: []*datapb.Binlog{deltalog},
|
|
}, nil
|
|
}
|