mirror of
https://gitee.com/milvus-io/milvus.git
synced 2026-01-07 19:31:51 +08:00
enhance: API integration with storage v2 in mix-compactions (#40008)
See #39173 --------- Signed-off-by: Ted Xu <ted.xu@zilliz.com>
This commit is contained in:
Ted Xu
GitHub
parent
ad36347fb3
commit
8562a102ec
@ -336,6 +336,7 @@ func (t *mixCompactionTask) BuildCompactionRequest() (*datapb.CompactionPlan, er
|
||||
Field2StatslogPaths: segInfo.GetStatslogs(),
|
||||
Deltalogs: segInfo.GetDeltalogs(),
|
||||
IsSorted: segInfo.GetIsSorted(),
|
||||
StorageVersion: segInfo.GetStorageVersion(),
|
||||
})
|
||||
segIDMap[segID] = segInfo.GetDeltalogs()
|
||||
}
|
||||
|
||||
@ -1666,7 +1666,7 @@ func (m *meta) completeMixCompactionMutation(t *datapb.CompactionTask, result *d
|
||||
CompactionFrom: compactFromSegIDs,
|
||||
LastExpireTime: tsoutil.ComposeTSByTime(time.Unix(t.GetStartTime(), 0), 0),
|
||||
Level: datapb.SegmentLevel_L1,
|
||||
|
||||
StorageVersion: compactToSegment.GetStorageVersion(),
|
||||
StartPosition: getMinPosition(lo.Map(compactFromSegInfos, func(info *SegmentInfo, _ int) *msgpb.MsgPosition {
|
||||
return info.GetStartPosition()
|
||||
})),
|
||||
|
||||
@ -7,7 +7,6 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/apache/arrow/go/v12/arrow/array"
|
||||
"github.com/samber/lo"
|
||||
"go.opentelemetry.io/otel"
|
||||
"go.uber.org/zap"
|
||||
|
||||
@ -23,24 +22,6 @@ import (
|
||||
"github.com/milvus-io/milvus/pkg/util/typeutil"
|
||||
)
|
||||
|
||||
type segmentWriterWrapper struct {
|
||||
*MultiSegmentWriter
|
||||
}
|
||||
|
||||
var _ storage.RecordWriter = (*segmentWriterWrapper)(nil)
|
||||
|
||||
func (w *segmentWriterWrapper) GetWrittenUncompressed() uint64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
func (w *segmentWriterWrapper) Write(record storage.Record) error {
|
||||
return w.MultiSegmentWriter.WriteRecord(record)
|
||||
}
|
||||
|
||||
func (w *segmentWriterWrapper) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func mergeSortMultipleSegments(ctx context.Context,
|
||||
plan *datapb.CompactionPlan,
|
||||
collectionID, partitionID, maxRows int64,
|
||||
@ -61,8 +42,7 @@ func mergeSortMultipleSegments(ctx context.Context,
|
||||
segIDAlloc := allocator.NewLocalAllocator(plan.GetPreAllocatedSegmentIDs().GetBegin(), plan.GetPreAllocatedSegmentIDs().GetEnd())
|
||||
logIDAlloc := allocator.NewLocalAllocator(plan.GetBeginLogID(), math.MaxInt64)
|
||||
compAlloc := NewCompactionAllocator(segIDAlloc, logIDAlloc)
|
||||
mWriter := NewMultiSegmentWriter(binlogIO, compAlloc, plan, maxRows, partitionID, collectionID, bm25FieldIds)
|
||||
writer := &segmentWriterWrapper{MultiSegmentWriter: mWriter}
|
||||
writer := NewMultiSegmentWriter(binlogIO, compAlloc, plan, maxRows, partitionID, collectionID, bm25FieldIds)
|
||||
|
||||
pkField, err := typeutil.GetPrimaryFieldSchema(plan.GetSchema())
|
||||
if err != nil {
|
||||
@ -126,16 +106,16 @@ func mergeSortMultipleSegments(ctx context.Context,
|
||||
log.Warn("compaction only support int64 and varchar pk field")
|
||||
}
|
||||
|
||||
if _, err = storage.MergeSort(plan.GetSchema(), segmentReaders, pkField.FieldID, writer, predicate); err != nil {
|
||||
if _, err = storage.MergeSort(plan.GetSchema(), segmentReaders, writer, predicate); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
res, err := mWriter.Finish()
|
||||
if err != nil {
|
||||
if err := writer.Close(); err != nil {
|
||||
log.Warn("compact wrong, failed to finish writer", zap.Error(err))
|
||||
return nil, err
|
||||
}
|
||||
|
||||
res := writer.GetCompactionSegments()
|
||||
for _, seg := range res {
|
||||
seg.IsSorted = true
|
||||
}
|
||||
@ -156,7 +136,6 @@ func mergeSortMultipleSegments(ctx context.Context,
|
||||
|
||||
totalElapse := tr.RecordSpan()
|
||||
log.Info("compact mergeSortMultipleSegments end",
|
||||
zap.Int64s("mergeSplit to segments", lo.Keys(mWriter.cachedMeta)),
|
||||
zap.Int("deleted row count", deletedRowCount),
|
||||
zap.Int("expired entities", expiredRowCount),
|
||||
zap.Int("missing deletes", missingDeleteCount),
|
||||
|
||||
@ -161,15 +161,14 @@ func (t *mixCompactionTask) mergeSplit(
|
||||
deletedRowCount += del
|
||||
expiredRowCount += exp
|
||||
}
|
||||
res, err := mWriter.Finish()
|
||||
if err != nil {
|
||||
if err := mWriter.Close(); err != nil {
|
||||
log.Warn("compact wrong, failed to finish writer", zap.Error(err))
|
||||
return nil, err
|
||||
}
|
||||
res := mWriter.GetCompactionSegments()
|
||||
|
||||
totalElapse := t.tr.RecordSpan()
|
||||
log.Info("compact mergeSplit end",
|
||||
zap.Int64s("mergeSplit to segments", lo.Keys(mWriter.cachedMeta)),
|
||||
zap.Int64("deleted row count", deletedRowCount),
|
||||
zap.Int64("expired entities", expiredRowCount),
|
||||
zap.Duration("total elapse", totalElapse))
|
||||
@ -194,30 +193,7 @@ func (t *mixCompactionTask) writeSegment(ctx context.Context,
|
||||
}
|
||||
entityFilter := newEntityFilter(delta, t.plan.GetCollectionTtl(), t.currentTime)
|
||||
|
||||
itr := 0
|
||||
binlogs := seg.GetFieldBinlogs()
|
||||
reader, err := storage.NewCompositeBinlogRecordReader(t.plan.GetSchema(), func() ([]*storage.Blob, error) {
|
||||
if len(binlogs) <= 0 {
|
||||
return nil, sio.EOF
|
||||
}
|
||||
paths := make([]string, len(binlogs))
|
||||
for i, fieldBinlog := range binlogs {
|
||||
if itr >= len(fieldBinlog.GetBinlogs()) {
|
||||
return nil, sio.EOF
|
||||
}
|
||||
paths[i] = fieldBinlog.GetBinlogs()[itr].GetLogPath()
|
||||
}
|
||||
itr++
|
||||
values, err := t.binlogIO.Download(ctx, paths)
|
||||
if err != nil {
|
||||
log.Warn("compact wrong, fail to download insertLogs", zap.Error(err))
|
||||
return nil, err
|
||||
}
|
||||
blobs := lo.Map(values, func(v []byte, i int) *storage.Blob {
|
||||
return &storage.Blob{Key: paths[i], Value: v}
|
||||
})
|
||||
return blobs, nil
|
||||
})
|
||||
reader, err := storage.NewBinlogRecordReader(ctx, seg.GetFieldBinlogs(), t.plan.GetSchema(), storage.WithDownloader(t.binlogIO.Download))
|
||||
if err != nil {
|
||||
log.Warn("compact wrong, failed to new insert binlogs reader", zap.Error(err))
|
||||
return
|
||||
@ -227,7 +203,7 @@ func (t *mixCompactionTask) writeSegment(ctx context.Context,
|
||||
writeSlice := func(r storage.Record, start, end int) error {
|
||||
sliced := r.Slice(start, end)
|
||||
defer sliced.Release()
|
||||
err = mWriter.WriteRecord(sliced)
|
||||
err = mWriter.Write(sliced)
|
||||
if err != nil {
|
||||
log.Warn("compact wrong, failed to writer row", zap.Error(err))
|
||||
return err
|
||||
@ -235,7 +211,8 @@ func (t *mixCompactionTask) writeSegment(ctx context.Context,
|
||||
return nil
|
||||
}
|
||||
for {
|
||||
err = reader.Next()
|
||||
var r storage.Record
|
||||
r, err = reader.Next()
|
||||
if err != nil {
|
||||
if err == sio.EOF {
|
||||
err = nil
|
||||
@ -245,7 +222,6 @@ func (t *mixCompactionTask) writeSegment(ctx context.Context,
|
||||
return
|
||||
}
|
||||
}
|
||||
r := reader.Record()
|
||||
pkArray := r.Column(pkField.FieldID)
|
||||
tsArray := r.Column(common.TimeStampField).(*array.Int64)
|
||||
|
||||
|
||||
@ -453,6 +453,7 @@ func (s *MixCompactionTaskSuite) TestSplitMergeEntityExpired() {
|
||||
}
|
||||
|
||||
func (s *MixCompactionTaskSuite) TestMergeNoExpirationLackBinlog() {
|
||||
s.T().Skip() // Skip added field related tests for now.
|
||||
s.initSegBuffer(1, 4)
|
||||
deleteTs := tsoutil.ComposeTSByTime(getMilvusBirthday().Add(10*time.Second), 0)
|
||||
tests := []struct {
|
||||
|
||||
@ -33,6 +33,7 @@ import (
|
||||
"github.com/milvus-io/milvus/internal/allocator"
|
||||
"github.com/milvus-io/milvus/internal/flushcommon/io"
|
||||
"github.com/milvus-io/milvus/internal/flushcommon/writebuffer"
|
||||
"github.com/milvus-io/milvus/internal/metastore/kv/binlog"
|
||||
"github.com/milvus-io/milvus/internal/storage"
|
||||
"github.com/milvus-io/milvus/pkg/common"
|
||||
"github.com/milvus-io/milvus/pkg/log"
|
||||
@ -47,8 +48,8 @@ type MultiSegmentWriter struct {
|
||||
binlogIO io.BinlogIO
|
||||
allocator *compactionAlloactor
|
||||
|
||||
writers []*SegmentWriter
|
||||
current int
|
||||
writer storage.BinlogRecordWriter
|
||||
currentSegmentID typeutil.UniqueID
|
||||
|
||||
maxRows int64
|
||||
segmentSize int64
|
||||
@ -61,14 +62,13 @@ type MultiSegmentWriter struct {
|
||||
collectionID int64
|
||||
channel string
|
||||
|
||||
cachedMeta map[typeutil.UniqueID]map[typeutil.UniqueID]*datapb.FieldBinlog
|
||||
// segID -> fieldID -> binlogs
|
||||
|
||||
res []*datapb.CompactionSegment
|
||||
// DONOT leave it empty of all segments are deleted, just return a segment with zero meta for datacoord
|
||||
bm25Fields []int64
|
||||
}
|
||||
|
||||
var _ storage.RecordWriter = &MultiSegmentWriter{}
|
||||
|
||||
type compactionAlloactor struct {
|
||||
segmentAlloc allocator.Interface
|
||||
logIDAlloc allocator.Interface
|
||||
@ -89,14 +89,13 @@ func (alloc *compactionAlloactor) getLogIDAllocator() allocator.Interface {
|
||||
return alloc.logIDAlloc
|
||||
}
|
||||
|
||||
func NewMultiSegmentWriter(binlogIO io.BinlogIO, allocator *compactionAlloactor, plan *datapb.CompactionPlan, maxRows int64, partitionID, collectionID int64, bm25Fields []int64) *MultiSegmentWriter {
|
||||
func NewMultiSegmentWriter(binlogIO io.BinlogIO, allocator *compactionAlloactor, plan *datapb.CompactionPlan,
|
||||
maxRows int64, partitionID, collectionID int64, bm25Fields []int64,
|
||||
) *MultiSegmentWriter {
|
||||
return &MultiSegmentWriter{
|
||||
binlogIO: binlogIO,
|
||||
allocator: allocator,
|
||||
|
||||
writers: make([]*SegmentWriter, 0),
|
||||
current: -1,
|
||||
|
||||
maxRows: maxRows, // For bloomfilter only
|
||||
segmentSize: plan.GetMaxSize(),
|
||||
|
||||
@ -105,177 +104,103 @@ func NewMultiSegmentWriter(binlogIO io.BinlogIO, allocator *compactionAlloactor,
|
||||
collectionID: collectionID,
|
||||
channel: plan.GetChannel(),
|
||||
|
||||
cachedMeta: make(map[typeutil.UniqueID]map[typeutil.UniqueID]*datapb.FieldBinlog),
|
||||
res: make([]*datapb.CompactionSegment, 0),
|
||||
bm25Fields: bm25Fields,
|
||||
}
|
||||
}
|
||||
|
||||
func (w *MultiSegmentWriter) finishCurrent() error {
|
||||
writer := w.writers[w.current]
|
||||
allBinlogs, ok := w.cachedMeta[writer.segmentID]
|
||||
if !ok {
|
||||
allBinlogs = make(map[typeutil.UniqueID]*datapb.FieldBinlog)
|
||||
}
|
||||
|
||||
if !writer.FlushAndIsEmpty() {
|
||||
kvs, partialBinlogs, err := serializeWrite(context.TODO(), w.allocator.getLogIDAllocator(), writer)
|
||||
if err != nil {
|
||||
func (w *MultiSegmentWriter) closeWriter() error {
|
||||
if w.writer != nil {
|
||||
if err := w.writer.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := w.binlogIO.Upload(context.TODO(), kvs); err != nil {
|
||||
return err
|
||||
fieldBinlogs, statsLog, bm25Logs := w.writer.GetLogs()
|
||||
|
||||
result := &datapb.CompactionSegment{
|
||||
SegmentID: w.currentSegmentID,
|
||||
InsertLogs: lo.Values(fieldBinlogs),
|
||||
Field2StatslogPaths: []*datapb.FieldBinlog{statsLog},
|
||||
NumOfRows: w.writer.GetRowNum(),
|
||||
Channel: w.channel,
|
||||
Bm25Logs: lo.Values(bm25Logs),
|
||||
}
|
||||
|
||||
mergeFieldBinlogs(allBinlogs, partialBinlogs)
|
||||
w.res = append(w.res, result)
|
||||
|
||||
log.Info("Segment writer flushed a segment",
|
||||
zap.Int64("segmentID", w.currentSegmentID),
|
||||
zap.String("channel", w.channel),
|
||||
zap.Int64("totalRows", w.writer.GetRowNum()))
|
||||
}
|
||||
|
||||
sPath, err := statSerializeWrite(context.TODO(), w.binlogIO, w.allocator.getLogIDAllocator(), writer)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
result := &datapb.CompactionSegment{
|
||||
SegmentID: writer.GetSegmentID(),
|
||||
InsertLogs: lo.Values(allBinlogs),
|
||||
Field2StatslogPaths: []*datapb.FieldBinlog{sPath},
|
||||
NumOfRows: writer.GetRowNum(),
|
||||
Channel: w.channel,
|
||||
}
|
||||
|
||||
if len(w.bm25Fields) > 0 {
|
||||
bmBinlogs, err := bm25SerializeWrite(context.TODO(), w.binlogIO, w.allocator.getLogIDAllocator(), writer)
|
||||
if err != nil {
|
||||
log.Warn("compact wrong, failed to serialize write segment bm25 stats", zap.Error(err))
|
||||
return err
|
||||
}
|
||||
result.Bm25Logs = bmBinlogs
|
||||
}
|
||||
|
||||
w.res = append(w.res, result)
|
||||
|
||||
log.Info("Segment writer flushed a segment",
|
||||
zap.Int64("segmentID", writer.GetSegmentID()),
|
||||
zap.String("channel", w.channel),
|
||||
zap.Int64("totalRows", writer.GetRowNum()),
|
||||
zap.Int64("totalSize", writer.GetTotalSize()))
|
||||
|
||||
w.cachedMeta[writer.segmentID] = nil
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *MultiSegmentWriter) addNewWriter() error {
|
||||
func (w *MultiSegmentWriter) rotateWriter() error {
|
||||
if err := w.closeWriter(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
newSegmentID, err := w.allocator.allocSegmentID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
writer, err := NewSegmentWriter(w.schema, w.maxRows, compactionBatchSize, newSegmentID, w.partitionID, w.collectionID, w.bm25Fields)
|
||||
w.currentSegmentID = newSegmentID
|
||||
|
||||
ctx := context.TODO()
|
||||
chunkSize := paramtable.Get().DataNodeCfg.BinLogMaxSize.GetAsUint64()
|
||||
rootPath := binlog.GetRootPath()
|
||||
|
||||
writer, err := storage.NewBinlogRecordWriter(ctx, w.collectionID, w.partitionID, newSegmentID,
|
||||
w.schema, w.allocator.logIDAlloc, chunkSize, rootPath, w.maxRows,
|
||||
storage.WithUploader(func(ctx context.Context, kvs map[string][]byte) error {
|
||||
return w.binlogIO.Upload(ctx, kvs)
|
||||
}))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
w.writers = append(w.writers, writer)
|
||||
w.current++
|
||||
w.writer = writer
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *MultiSegmentWriter) getWriter() (*SegmentWriter, error) {
|
||||
if len(w.writers) == 0 {
|
||||
if err := w.addNewWriter(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return w.writers[w.current], nil
|
||||
func (w *MultiSegmentWriter) GetWrittenUncompressed() uint64 {
|
||||
if w.writer == nil {
|
||||
return 0
|
||||
}
|
||||
|
||||
if w.writers[w.current].GetTotalSize() > w.segmentSize {
|
||||
if err := w.finishCurrent(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := w.addNewWriter(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return w.writers[w.current], nil
|
||||
return w.writer.GetWrittenUncompressed()
|
||||
}
|
||||
|
||||
func (w *MultiSegmentWriter) writeInternal(writer *SegmentWriter) error {
|
||||
if writer.IsFull() {
|
||||
// init segment fieldBinlogs if it is not exist
|
||||
if _, ok := w.cachedMeta[writer.segmentID]; !ok {
|
||||
w.cachedMeta[writer.segmentID] = make(map[typeutil.UniqueID]*datapb.FieldBinlog)
|
||||
}
|
||||
func (w *MultiSegmentWriter) GetCompactionSegments() []*datapb.CompactionSegment {
|
||||
return w.res
|
||||
}
|
||||
|
||||
kvs, partialBinlogs, err := serializeWrite(context.TODO(), w.allocator.getLogIDAllocator(), writer)
|
||||
if err != nil {
|
||||
func (w *MultiSegmentWriter) Write(r storage.Record) error {
|
||||
if w.writer == nil || w.writer.GetWrittenUncompressed() >= uint64(w.segmentSize) {
|
||||
if err := w.rotateWriter(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := w.binlogIO.Upload(context.TODO(), kvs); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
mergeFieldBinlogs(w.cachedMeta[writer.segmentID], partialBinlogs)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *MultiSegmentWriter) WriteRecord(r storage.Record) error {
|
||||
writer, err := w.getWriter()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := w.writeInternal(writer); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return writer.WriteRecord(r)
|
||||
}
|
||||
|
||||
func (w *MultiSegmentWriter) Write(v *storage.Value) error {
|
||||
writer, err := w.getWriter()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := w.writeInternal(writer); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return writer.Write(v)
|
||||
}
|
||||
|
||||
func (w *MultiSegmentWriter) appendEmptySegment() error {
|
||||
writer, err := w.getWriter()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
w.res = append(w.res, &datapb.CompactionSegment{
|
||||
SegmentID: writer.GetSegmentID(),
|
||||
NumOfRows: 0,
|
||||
Channel: w.channel,
|
||||
})
|
||||
return nil
|
||||
return w.writer.Write(r)
|
||||
}
|
||||
|
||||
// DONOT return an empty list if every insert of the segment is deleted,
|
||||
// append an empty segment instead
|
||||
func (w *MultiSegmentWriter) Finish() ([]*datapb.CompactionSegment, error) {
|
||||
if w.current == -1 {
|
||||
if err := w.appendEmptySegment(); err != nil {
|
||||
return nil, err
|
||||
func (w *MultiSegmentWriter) Close() error {
|
||||
if w.writer == nil && len(w.res) == 0 {
|
||||
// append an empty segment
|
||||
id, err := w.allocator.segmentAlloc.AllocOne()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return w.res, nil
|
||||
w.res = append(w.res, &datapb.CompactionSegment{
|
||||
SegmentID: id,
|
||||
NumOfRows: 0,
|
||||
Channel: w.channel,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
if !w.writers[w.current].FlushAndIsEmpty() {
|
||||
if err := w.finishCurrent(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
return w.res, nil
|
||||
return w.closeWriter()
|
||||
}
|
||||
|
||||
func NewSegmentDeltaWriter(segmentID, partitionID, collectionID int64) *SegmentDeltaWriter {
|
||||
|
||||
@ -22,6 +22,7 @@ import (
|
||||
"github.com/stretchr/testify/suite"
|
||||
|
||||
"github.com/milvus-io/milvus/internal/storage"
|
||||
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||||
"github.com/milvus-io/milvus/pkg/util/tsoutil"
|
||||
)
|
||||
|
||||
@ -41,6 +42,7 @@ func (s *SegmentWriteSuite) SetupSuite() {
|
||||
}
|
||||
|
||||
func (s *SegmentWriteSuite) TestWriteFailed() {
|
||||
paramtable.Init()
|
||||
s.Run("get bm25 field failed", func() {
|
||||
schema := genCollectionSchemaWithBM25()
|
||||
// init segment writer with invalid bm25 fieldID
|
||||
|
||||
@ -29,6 +29,7 @@ import (
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/internal/allocator"
|
||||
"github.com/milvus-io/milvus/internal/datanode/compaction"
|
||||
iter "github.com/milvus-io/milvus/internal/datanode/iterators"
|
||||
"github.com/milvus-io/milvus/internal/flushcommon/io"
|
||||
@ -41,7 +42,6 @@ import (
|
||||
"github.com/milvus-io/milvus/pkg/proto/indexcgopb"
|
||||
"github.com/milvus-io/milvus/pkg/proto/indexpb"
|
||||
"github.com/milvus-io/milvus/pkg/proto/workerpb"
|
||||
"github.com/milvus-io/milvus/pkg/util/conc"
|
||||
_ "github.com/milvus-io/milvus/pkg/util/funcutil"
|
||||
"github.com/milvus-io/milvus/pkg/util/metautil"
|
||||
"github.com/milvus-io/milvus/pkg/util/timerecord"
|
||||
@ -64,7 +64,6 @@ type statsTask struct {
|
||||
node *IndexNode
|
||||
binlogIO io.BinlogIO
|
||||
|
||||
insertLogs [][]string
|
||||
deltaLogs []string
|
||||
logIDOffset int64
|
||||
}
|
||||
@ -137,16 +136,6 @@ func (st *statsTask) PreExecute(ctx context.Context) error {
|
||||
return err
|
||||
}
|
||||
|
||||
st.insertLogs = make([][]string, 0)
|
||||
binlogNum := len(st.req.GetInsertLogs()[0].GetBinlogs())
|
||||
for idx := 0; idx < binlogNum; idx++ {
|
||||
var batchPaths []string
|
||||
for _, f := range st.req.GetInsertLogs() {
|
||||
batchPaths = append(batchPaths, f.GetBinlogs()[idx].GetLogPath())
|
||||
}
|
||||
st.insertLogs = append(st.insertLogs, batchPaths)
|
||||
}
|
||||
|
||||
for _, d := range st.req.GetDeltaLogs() {
|
||||
for _, l := range d.GetBinlogs() {
|
||||
st.deltaLogs = append(st.deltaLogs, l.GetLogPath())
|
||||
@ -156,205 +145,32 @@ func (st *statsTask) PreExecute(ctx context.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// segmentRecordWriter is a wrapper of SegmentWriter to implement RecordWriter interface
|
||||
type segmentRecordWriter struct {
|
||||
sw *compaction.SegmentWriter
|
||||
binlogMaxSize uint64
|
||||
rootPath string
|
||||
logID int64
|
||||
maxLogID int64
|
||||
binlogIO io.BinlogIO
|
||||
ctx context.Context
|
||||
numRows int64
|
||||
bm25FieldIds []int64
|
||||
|
||||
lastUploads []*conc.Future[any]
|
||||
binlogs map[typeutil.UniqueID]*datapb.FieldBinlog
|
||||
statslog *datapb.FieldBinlog
|
||||
bm25statslog []*datapb.FieldBinlog
|
||||
}
|
||||
|
||||
var _ storage.RecordWriter = (*segmentRecordWriter)(nil)
|
||||
|
||||
func (srw *segmentRecordWriter) Close() error {
|
||||
if !srw.sw.FlushAndIsEmpty() {
|
||||
if err := srw.upload(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := srw.waitLastUpload(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
statslog, err := srw.statSerializeWrite()
|
||||
if err != nil {
|
||||
log.Ctx(srw.ctx).Warn("stats wrong, failed to serialize write segment stats",
|
||||
zap.Int64("remaining row count", srw.numRows), zap.Error(err))
|
||||
return err
|
||||
}
|
||||
srw.statslog = statslog
|
||||
srw.logID++
|
||||
|
||||
if len(srw.bm25FieldIds) > 0 {
|
||||
binlogNums, bm25StatsLogs, err := srw.bm25SerializeWrite()
|
||||
if err != nil {
|
||||
log.Ctx(srw.ctx).Warn("compact wrong, failed to serialize write segment bm25 stats", zap.Error(err))
|
||||
return err
|
||||
}
|
||||
srw.logID += binlogNums
|
||||
srw.bm25statslog = bm25StatsLogs
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (srw *segmentRecordWriter) GetWrittenUncompressed() uint64 {
|
||||
return srw.sw.WrittenMemorySize()
|
||||
}
|
||||
|
||||
func (srw *segmentRecordWriter) Write(r storage.Record) error {
|
||||
err := srw.sw.WriteRecord(r)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if srw.sw.IsFullWithBinlogMaxSize(srw.binlogMaxSize) {
|
||||
return srw.upload()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (srw *segmentRecordWriter) upload() error {
|
||||
if err := srw.waitLastUpload(); err != nil {
|
||||
return err
|
||||
}
|
||||
binlogNum, kvs, partialBinlogs, err := serializeWrite(srw.ctx, srw.rootPath, srw.logID, srw.sw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
srw.lastUploads = srw.binlogIO.AsyncUpload(srw.ctx, kvs)
|
||||
if srw.binlogs == nil {
|
||||
srw.binlogs = make(map[typeutil.UniqueID]*datapb.FieldBinlog)
|
||||
}
|
||||
mergeFieldBinlogs(srw.binlogs, partialBinlogs)
|
||||
|
||||
srw.logID += binlogNum
|
||||
if srw.logID > srw.maxLogID {
|
||||
return fmt.Errorf("log id exausted")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (srw *segmentRecordWriter) waitLastUpload() error {
|
||||
if len(srw.lastUploads) > 0 {
|
||||
for _, future := range srw.lastUploads {
|
||||
if _, err := future.Await(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (srw *segmentRecordWriter) statSerializeWrite() (*datapb.FieldBinlog, error) {
|
||||
ctx, span := otel.Tracer(typeutil.DataNodeRole).Start(srw.ctx, "statslog serializeWrite")
|
||||
defer span.End()
|
||||
sblob, err := srw.sw.Finish()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
key, _ := binlog.BuildLogPathWithRootPath(srw.rootPath, storage.StatsBinlog,
|
||||
srw.sw.GetCollectionID(), srw.sw.GetPartitionID(), srw.sw.GetSegmentID(), srw.sw.GetPkID(), srw.logID)
|
||||
kvs := map[string][]byte{key: sblob.GetValue()}
|
||||
statFieldLog := &datapb.FieldBinlog{
|
||||
FieldID: srw.sw.GetPkID(),
|
||||
Binlogs: []*datapb.Binlog{
|
||||
{
|
||||
LogSize: int64(len(sblob.GetValue())),
|
||||
MemorySize: int64(len(sblob.GetValue())),
|
||||
LogPath: key,
|
||||
EntriesNum: srw.numRows,
|
||||
},
|
||||
},
|
||||
}
|
||||
if err := srw.binlogIO.Upload(ctx, kvs); err != nil {
|
||||
log.Ctx(ctx).Warn("failed to upload insert log", zap.Error(err))
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return statFieldLog, nil
|
||||
}
|
||||
|
||||
func (srw *segmentRecordWriter) bm25SerializeWrite() (int64, []*datapb.FieldBinlog, error) {
|
||||
ctx, span := otel.Tracer(typeutil.DataNodeRole).Start(srw.ctx, "bm25log serializeWrite")
|
||||
defer span.End()
|
||||
writer := srw.sw
|
||||
stats, err := writer.GetBm25StatsBlob()
|
||||
if err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
|
||||
kvs := make(map[string][]byte)
|
||||
binlogs := []*datapb.FieldBinlog{}
|
||||
cnt := int64(0)
|
||||
for fieldID, blob := range stats {
|
||||
key, _ := binlog.BuildLogPathWithRootPath(srw.rootPath, storage.BM25Binlog,
|
||||
writer.GetCollectionID(), writer.GetPartitionID(), writer.GetSegmentID(), fieldID, srw.logID)
|
||||
kvs[key] = blob.GetValue()
|
||||
fieldLog := &datapb.FieldBinlog{
|
||||
FieldID: fieldID,
|
||||
Binlogs: []*datapb.Binlog{
|
||||
{
|
||||
LogSize: int64(len(blob.GetValue())),
|
||||
MemorySize: int64(len(blob.GetValue())),
|
||||
LogPath: key,
|
||||
EntriesNum: srw.numRows,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
binlogs = append(binlogs, fieldLog)
|
||||
srw.logID++
|
||||
cnt++
|
||||
}
|
||||
|
||||
if err := srw.binlogIO.Upload(ctx, kvs); err != nil {
|
||||
log.Ctx(ctx).Warn("failed to upload bm25 log", zap.Error(err))
|
||||
return 0, nil, err
|
||||
}
|
||||
|
||||
return cnt, binlogs, nil
|
||||
}
|
||||
|
||||
func (st *statsTask) sort(ctx context.Context) ([]*datapb.FieldBinlog, error) {
|
||||
numRows := st.req.GetNumRows()
|
||||
|
||||
bm25FieldIds := compaction.GetBM25FieldIDs(st.req.GetSchema())
|
||||
pkField, err := typeutil.GetPrimaryFieldSchema(st.req.GetSchema())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pkFieldID := pkField.FieldID
|
||||
writer, err := compaction.NewSegmentWriter(st.req.GetSchema(), numRows, statsBatchSize,
|
||||
st.req.GetTargetSegmentID(), st.req.GetPartitionID(), st.req.GetCollectionID(), bm25FieldIds)
|
||||
|
||||
alloc := allocator.NewLocalAllocator(st.req.StartLogID, st.req.EndLogID)
|
||||
srw, err := storage.NewBinlogRecordWriter(ctx,
|
||||
st.req.GetCollectionID(),
|
||||
st.req.GetPartitionID(),
|
||||
st.req.GetTargetSegmentID(),
|
||||
st.req.GetSchema(),
|
||||
alloc,
|
||||
st.req.GetBinlogMaxSize(),
|
||||
st.req.GetStorageConfig().RootPath,
|
||||
numRows,
|
||||
storage.WithUploader(func(ctx context.Context, kvs map[string][]byte) error {
|
||||
return st.binlogIO.Upload(ctx, kvs)
|
||||
}))
|
||||
if err != nil {
|
||||
log.Ctx(ctx).Warn("sort segment wrong, unable to init segment writer",
|
||||
zap.Int64("taskID", st.req.GetTaskID()), zap.Error(err))
|
||||
return nil, err
|
||||
}
|
||||
srw := &segmentRecordWriter{
|
||||
sw: writer,
|
||||
binlogMaxSize: st.req.GetBinlogMaxSize(),
|
||||
rootPath: st.req.GetStorageConfig().GetRootPath(),
|
||||
logID: st.req.StartLogID,
|
||||
maxLogID: st.req.EndLogID,
|
||||
binlogIO: st.binlogIO,
|
||||
ctx: ctx,
|
||||
numRows: st.req.NumRows,
|
||||
bm25FieldIds: bm25FieldIds,
|
||||
}
|
||||
|
||||
log := log.Ctx(ctx).With(
|
||||
zap.String("clusterID", st.req.GetClusterID()),
|
||||
@ -362,7 +178,6 @@ func (st *statsTask) sort(ctx context.Context) ([]*datapb.FieldBinlog, error) {
|
||||
zap.Int64("collectionID", st.req.GetCollectionID()),
|
||||
zap.Int64("partitionID", st.req.GetPartitionID()),
|
||||
zap.Int64("segmentID", st.req.GetSegmentID()),
|
||||
zap.Int64s("bm25Fields", bm25FieldIds),
|
||||
)
|
||||
|
||||
deletePKs, err := st.loadDeltalogs(ctx, st.deltaLogs)
|
||||
@ -395,38 +210,13 @@ func (st *statsTask) sort(ctx context.Context) ([]*datapb.FieldBinlog, error) {
|
||||
}
|
||||
}
|
||||
|
||||
downloadTimeCost := time.Duration(0)
|
||||
|
||||
rrs := make([]storage.RecordReader, len(st.insertLogs))
|
||||
|
||||
for i, paths := range st.insertLogs {
|
||||
log := log.With(zap.Strings("paths", paths))
|
||||
downloadStart := time.Now()
|
||||
allValues, err := st.binlogIO.Download(ctx, paths)
|
||||
if err != nil {
|
||||
log.Warn("download wrong, fail to download insertLogs", zap.Error(err))
|
||||
return nil, err
|
||||
}
|
||||
downloadTimeCost += time.Since(downloadStart)
|
||||
|
||||
blobs := lo.Map(allValues, func(v []byte, i int) *storage.Blob {
|
||||
return &storage.Blob{Key: paths[i], Value: v}
|
||||
})
|
||||
|
||||
rr, err := storage.NewCompositeBinlogRecordReader(st.req.Schema, storage.MakeBlobsReader(blobs))
|
||||
if err != nil {
|
||||
log.Warn("downloadData wrong, failed to new insert binlogs reader", zap.Error(err))
|
||||
return nil, err
|
||||
}
|
||||
rrs[i] = rr
|
||||
rr, err := storage.NewBinlogRecordReader(ctx, st.req.InsertLogs, st.req.Schema, storage.WithDownloader(st.binlogIO.Download))
|
||||
if err != nil {
|
||||
log.Warn("error creating insert binlog reader", zap.Error(err))
|
||||
return nil, err
|
||||
}
|
||||
|
||||
log.Info("download data success",
|
||||
zap.Int64("numRows", numRows),
|
||||
zap.Duration("download binlogs elapse", downloadTimeCost),
|
||||
)
|
||||
|
||||
numValidRows, err := storage.Sort(st.req.Schema, rrs, writer.GetPkID(), srw, isValueValid)
|
||||
rrs := []storage.RecordReader{rr}
|
||||
numValidRows, err := storage.Sort(st.req.Schema, rrs, srw, isValueValid)
|
||||
if err != nil {
|
||||
log.Warn("sort failed", zap.Int64("taskID", st.req.GetTaskID()), zap.Error(err))
|
||||
return nil, err
|
||||
@ -435,17 +225,18 @@ func (st *statsTask) sort(ctx context.Context) ([]*datapb.FieldBinlog, error) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
insertLogs := lo.Values(srw.binlogs)
|
||||
binlogs, stats, bm25stats := srw.GetLogs()
|
||||
insertLogs := lo.Values(binlogs)
|
||||
if err := binlog.CompressFieldBinlogs(insertLogs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
statsLogs := []*datapb.FieldBinlog{srw.statslog}
|
||||
statsLogs := []*datapb.FieldBinlog{stats}
|
||||
if err := binlog.CompressFieldBinlogs(statsLogs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
bm25StatsLogs := srw.bm25statslog
|
||||
bm25StatsLogs := lo.Values(bm25stats)
|
||||
if err := binlog.CompressFieldBinlogs(bm25StatsLogs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@ -32,6 +32,7 @@ import (
|
||||
"github.com/milvus-io/milvus/internal/flushcommon/io"
|
||||
"github.com/milvus-io/milvus/internal/storage"
|
||||
"github.com/milvus-io/milvus/pkg/common"
|
||||
"github.com/milvus-io/milvus/pkg/proto/indexpb"
|
||||
"github.com/milvus-io/milvus/pkg/proto/workerpb"
|
||||
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||||
"github.com/milvus-io/milvus/pkg/util/tsoutil"
|
||||
@ -95,7 +96,6 @@ func (s *TaskStatsSuite) TestSortSegmentWithBM25() {
|
||||
return result, nil
|
||||
})
|
||||
s.mockBinlogIO.EXPECT().Upload(mock.Anything, mock.Anything).Return(nil)
|
||||
s.mockBinlogIO.EXPECT().AsyncUpload(mock.Anything, mock.Anything).Return(nil)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
@ -111,8 +111,11 @@ func (s *TaskStatsSuite) TestSortSegmentWithBM25() {
|
||||
Schema: s.schema,
|
||||
NumRows: 1,
|
||||
StartLogID: 0,
|
||||
EndLogID: 5,
|
||||
EndLogID: 7,
|
||||
BinlogMaxSize: 64 * 1024 * 1024,
|
||||
StorageConfig: &indexpb.StorageConfig{
|
||||
RootPath: "root_path",
|
||||
},
|
||||
}, node, s.mockBinlogIO)
|
||||
err = task.PreExecute(ctx)
|
||||
s.Require().NoError(err)
|
||||
@ -142,7 +145,6 @@ func (s *TaskStatsSuite) TestSortSegmentWithBM25() {
|
||||
return result, nil
|
||||
})
|
||||
s.mockBinlogIO.EXPECT().Upload(mock.Anything, mock.Anything).Return(fmt.Errorf("mock error")).Once()
|
||||
s.mockBinlogIO.EXPECT().AsyncUpload(mock.Anything, mock.Anything).Return(nil)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
@ -158,8 +160,11 @@ func (s *TaskStatsSuite) TestSortSegmentWithBM25() {
|
||||
Schema: s.schema,
|
||||
NumRows: 1,
|
||||
StartLogID: 0,
|
||||
EndLogID: 5,
|
||||
EndLogID: 7,
|
||||
BinlogMaxSize: 64 * 1024 * 1024,
|
||||
StorageConfig: &indexpb.StorageConfig{
|
||||
RootPath: "root_path",
|
||||
},
|
||||
}, node, s.mockBinlogIO)
|
||||
err = task.PreExecute(ctx)
|
||||
s.Require().NoError(err)
|
||||
|
||||
@ -185,12 +185,16 @@ func DecompressBinLogWithRootPath(rootPath string, binlogType storage.BinlogType
|
||||
return nil
|
||||
}
|
||||
|
||||
func GetRootPath() string {
|
||||
if paramtable.Get().CommonCfg.StorageType.GetValue() == "local" {
|
||||
return paramtable.Get().LocalStorageCfg.Path.GetValue()
|
||||
}
|
||||
return paramtable.Get().MinioCfg.RootPath.GetValue()
|
||||
}
|
||||
|
||||
// build a binlog path on the storage by metadata
|
||||
func BuildLogPath(binlogType storage.BinlogType, collectionID, partitionID, segmentID, fieldID, logID typeutil.UniqueID) (string, error) {
|
||||
chunkManagerRootPath := paramtable.Get().MinioCfg.RootPath.GetValue()
|
||||
if paramtable.Get().CommonCfg.StorageType.GetValue() == "local" {
|
||||
chunkManagerRootPath = paramtable.Get().LocalStorageCfg.Path.GetValue()
|
||||
}
|
||||
chunkManagerRootPath := GetRootPath()
|
||||
return BuildLogPathWithRootPath(chunkManagerRootPath, binlogType, collectionID, partitionID, segmentID, fieldID, logID)
|
||||
}
|
||||
|
||||
|
||||
135
internal/storage/rw.go
Normal file
135
internal/storage/rw.go
Normal file
@ -0,0 +1,135 @@
|
||||
// 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 storage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
sio "io"
|
||||
|
||||
"github.com/samber/lo"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/internal/allocator"
|
||||
"github.com/milvus-io/milvus/pkg/proto/datapb"
|
||||
"github.com/milvus-io/milvus/pkg/util/merr"
|
||||
)
|
||||
|
||||
const (
|
||||
StorageV1 int64 = 0
|
||||
StorageV2 int64 = 2
|
||||
)
|
||||
|
||||
type rwOptions struct {
|
||||
version int64
|
||||
bufferSize uint64
|
||||
downloader func(ctx context.Context, paths []string) ([][]byte, error)
|
||||
uploader func(ctx context.Context, kvs map[string][]byte) error
|
||||
}
|
||||
|
||||
type RwOption func(*rwOptions)
|
||||
|
||||
func defaultRwOptions() *rwOptions {
|
||||
return &rwOptions{
|
||||
bufferSize: 32 * 1024 * 1024,
|
||||
}
|
||||
}
|
||||
|
||||
func WithVersion(version int64) RwOption {
|
||||
return func(options *rwOptions) {
|
||||
options.version = version
|
||||
}
|
||||
}
|
||||
|
||||
func WithBufferSize(bufferSize uint64) RwOption {
|
||||
return func(options *rwOptions) {
|
||||
options.bufferSize = bufferSize
|
||||
}
|
||||
}
|
||||
|
||||
func WithDownloader(downloader func(ctx context.Context, paths []string) ([][]byte, error)) RwOption {
|
||||
return func(options *rwOptions) {
|
||||
options.downloader = downloader
|
||||
}
|
||||
}
|
||||
|
||||
func WithUploader(uploader func(ctx context.Context, kvs map[string][]byte) error) RwOption {
|
||||
return func(options *rwOptions) {
|
||||
options.uploader = uploader
|
||||
}
|
||||
}
|
||||
|
||||
func NewBinlogRecordReader(ctx context.Context, binlogs []*datapb.FieldBinlog, schema *schemapb.CollectionSchema, option ...RwOption) (RecordReader, error) {
|
||||
rwOptions := defaultRwOptions()
|
||||
for _, opt := range option {
|
||||
opt(rwOptions)
|
||||
}
|
||||
switch rwOptions.version {
|
||||
case StorageV1:
|
||||
itr := 0
|
||||
return newCompositeBinlogRecordReader(schema, func() ([]*Blob, error) {
|
||||
if len(binlogs) <= 0 {
|
||||
return nil, sio.EOF
|
||||
}
|
||||
paths := make([]string, len(binlogs))
|
||||
for i, fieldBinlog := range binlogs {
|
||||
if itr >= len(fieldBinlog.GetBinlogs()) {
|
||||
return nil, sio.EOF
|
||||
}
|
||||
paths[i] = fieldBinlog.GetBinlogs()[itr].GetLogPath()
|
||||
}
|
||||
itr++
|
||||
values, err := rwOptions.downloader(ctx, paths)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
blobs := lo.Map(values, func(v []byte, i int) *Blob {
|
||||
return &Blob{Key: paths[i], Value: v}
|
||||
})
|
||||
return blobs, nil
|
||||
})
|
||||
case StorageV2:
|
||||
// TODO: integrate v2
|
||||
}
|
||||
return nil, merr.WrapErrServiceInternal(fmt.Sprintf("unsupported storage version %d", rwOptions.version))
|
||||
}
|
||||
|
||||
func NewBinlogRecordWriter(ctx context.Context, collectionID, partitionID, segmentID UniqueID,
|
||||
schema *schemapb.CollectionSchema, allocator allocator.Interface, chunkSize uint64, rootPath string, maxRowNum int64,
|
||||
option ...RwOption,
|
||||
) (BinlogRecordWriter, error) {
|
||||
rwOptions := defaultRwOptions()
|
||||
for _, opt := range option {
|
||||
opt(rwOptions)
|
||||
}
|
||||
switch rwOptions.version {
|
||||
case StorageV1:
|
||||
blobsWriter := func(blobs []*Blob) error {
|
||||
kvs := make(map[string][]byte, len(blobs))
|
||||
for _, blob := range blobs {
|
||||
kvs[blob.Key] = blob.Value
|
||||
}
|
||||
return rwOptions.uploader(ctx, kvs)
|
||||
}
|
||||
return newCompositeBinlogRecordWriter(collectionID, partitionID, segmentID, schema,
|
||||
blobsWriter, allocator, chunkSize, rootPath, maxRowNum,
|
||||
)
|
||||
case StorageV2:
|
||||
// TODO: integrate v2
|
||||
}
|
||||
return nil, merr.WrapErrServiceInternal(fmt.Sprintf("unsupported storage version %d", rwOptions.version))
|
||||
}
|
||||
@ -43,8 +43,7 @@ type Record interface {
|
||||
}
|
||||
|
||||
type RecordReader interface {
|
||||
Next() error
|
||||
Record() Record
|
||||
Next() (Record, error)
|
||||
Close() error
|
||||
}
|
||||
|
||||
@ -529,11 +528,12 @@ type DeserializeReader[T any] struct {
|
||||
// Iterate to next value, return error or EOF if no more value.
|
||||
func (deser *DeserializeReader[T]) Next() error {
|
||||
if deser.rec == nil || deser.pos >= deser.rec.Len()-1 {
|
||||
if err := deser.rr.Next(); err != nil {
|
||||
r, err := deser.rr.Next()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
deser.pos = 0
|
||||
deser.rec = deser.rr.Record()
|
||||
deser.rec = r
|
||||
|
||||
deser.values = make([]T, deser.rec.Len())
|
||||
|
||||
|
||||
@ -21,16 +21,22 @@ import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"sort"
|
||||
"strconv"
|
||||
|
||||
"github.com/apache/arrow/go/v12/arrow"
|
||||
"github.com/apache/arrow/go/v12/arrow/array"
|
||||
"github.com/apache/arrow/go/v12/arrow/memory"
|
||||
"github.com/samber/lo"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/internal/allocator"
|
||||
"github.com/milvus-io/milvus/internal/json"
|
||||
"github.com/milvus-io/milvus/pkg/common"
|
||||
"github.com/milvus-io/milvus/pkg/log"
|
||||
"github.com/milvus-io/milvus/pkg/proto/datapb"
|
||||
"github.com/milvus-io/milvus/pkg/proto/etcdpb"
|
||||
"github.com/milvus-io/milvus/pkg/util/merr"
|
||||
"github.com/milvus-io/milvus/pkg/util/metautil"
|
||||
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||||
@ -50,7 +56,6 @@ type CompositeBinlogRecordReader struct {
|
||||
|
||||
schema *schemapb.CollectionSchema
|
||||
index map[FieldID]int16
|
||||
r *compositeRecord
|
||||
}
|
||||
|
||||
func (crr *CompositeBinlogRecordReader) iterateNextBatch() error {
|
||||
@ -95,45 +100,45 @@ func (crr *CompositeBinlogRecordReader) iterateNextBatch() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (crr *CompositeBinlogRecordReader) Next() error {
|
||||
func (crr *CompositeBinlogRecordReader) Next() (Record, error) {
|
||||
if crr.rrs == nil {
|
||||
if err := crr.iterateNextBatch(); err != nil {
|
||||
return err
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
composeRecord := func() bool {
|
||||
composeRecord := func() (Record, bool) {
|
||||
recs := make([]arrow.Array, len(crr.rrs))
|
||||
for i, rr := range crr.rrs {
|
||||
if ok := rr.Next(); !ok {
|
||||
return false
|
||||
return nil, false
|
||||
}
|
||||
recs[i] = rr.Record().Column(0)
|
||||
}
|
||||
crr.r = &compositeRecord{
|
||||
return &compositeRecord{
|
||||
index: crr.index,
|
||||
recs: recs,
|
||||
}
|
||||
return true
|
||||
}, true
|
||||
}
|
||||
|
||||
// Try compose records
|
||||
if ok := composeRecord(); !ok {
|
||||
var (
|
||||
r Record
|
||||
ok bool
|
||||
)
|
||||
r, ok = composeRecord()
|
||||
if !ok {
|
||||
// If failed the first time, try iterate next batch (blob), the error may be io.EOF
|
||||
if err := crr.iterateNextBatch(); err != nil {
|
||||
return err
|
||||
return nil, err
|
||||
}
|
||||
// If iterate next batch success, try compose again
|
||||
if ok := composeRecord(); !ok {
|
||||
if r, ok = composeRecord(); !ok {
|
||||
// If the next blob is empty, return io.EOF (it's rare).
|
||||
return io.EOF
|
||||
return nil, io.EOF
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (crr *CompositeBinlogRecordReader) Record() Record {
|
||||
return crr.r
|
||||
return r, nil
|
||||
}
|
||||
|
||||
func (crr *CompositeBinlogRecordReader) Close() error {
|
||||
@ -151,7 +156,6 @@ func (crr *CompositeBinlogRecordReader) Close() error {
|
||||
}
|
||||
}
|
||||
}
|
||||
crr.r = nil
|
||||
return nil
|
||||
}
|
||||
|
||||
@ -200,7 +204,7 @@ func MakeBlobsReader(blobs []*Blob) ChunkedBlobsReader {
|
||||
}
|
||||
}
|
||||
|
||||
func NewCompositeBinlogRecordReader(schema *schemapb.CollectionSchema, blobsReader ChunkedBlobsReader) (*CompositeBinlogRecordReader, error) {
|
||||
func newCompositeBinlogRecordReader(schema *schemapb.CollectionSchema, blobsReader ChunkedBlobsReader) (*CompositeBinlogRecordReader, error) {
|
||||
return &CompositeBinlogRecordReader{
|
||||
schema: schema,
|
||||
BlobsReader: blobsReader,
|
||||
@ -272,7 +276,7 @@ func ValueDeserializer(r Record, v []*Value, fieldSchema []*schemapb.FieldSchema
|
||||
}
|
||||
|
||||
func NewBinlogDeserializeReader(schema *schemapb.CollectionSchema, blobsReader ChunkedBlobsReader) (*DeserializeReader[*Value], error) {
|
||||
reader, err := NewCompositeBinlogRecordReader(schema, blobsReader)
|
||||
reader, err := newCompositeBinlogRecordReader(schema, blobsReader)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@ -283,7 +287,7 @@ func NewBinlogDeserializeReader(schema *schemapb.CollectionSchema, blobsReader C
|
||||
}
|
||||
|
||||
func newDeltalogOneFieldReader(blobs []*Blob) (*DeserializeReader[*DeleteLog], error) {
|
||||
reader, err := NewCompositeBinlogRecordReader(nil, MakeBlobsReader(blobs))
|
||||
reader, err := newCompositeBinlogRecordReader(nil, MakeBlobsReader(blobs))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@ -439,6 +443,345 @@ func ValueSerializer(v []*Value, fieldSchema []*schemapb.FieldSchema) (Record, e
|
||||
return NewSimpleArrowRecord(array.NewRecord(arrow.NewSchema(fields, nil), arrays, int64(len(v))), field2Col), nil
|
||||
}
|
||||
|
||||
type BinlogRecordWriter interface {
|
||||
RecordWriter
|
||||
GetLogs() (
|
||||
fieldBinlogs map[FieldID]*datapb.FieldBinlog,
|
||||
statsLog *datapb.FieldBinlog,
|
||||
bm25StatsLog map[FieldID]*datapb.FieldBinlog,
|
||||
)
|
||||
GetRowNum() int64
|
||||
}
|
||||
|
||||
type ChunkedBlobsWriter func([]*Blob) error
|
||||
|
||||
type CompositeBinlogRecordWriter struct {
|
||||
// attributes
|
||||
collectionID UniqueID
|
||||
partitionID UniqueID
|
||||
segmentID UniqueID
|
||||
schema *schemapb.CollectionSchema
|
||||
BlobsWriter ChunkedBlobsWriter
|
||||
allocator allocator.Interface
|
||||
chunkSize uint64
|
||||
rootPath string
|
||||
maxRowNum int64
|
||||
pkstats *PrimaryKeyStats
|
||||
bm25Stats map[int64]*BM25Stats
|
||||
|
||||
// writers and stats generated at runtime
|
||||
fieldWriters map[FieldID]*BinlogStreamWriter
|
||||
rw RecordWriter
|
||||
tsFrom typeutil.Timestamp
|
||||
tsTo typeutil.Timestamp
|
||||
rowNum int64
|
||||
|
||||
// results
|
||||
fieldBinlogs map[FieldID]*datapb.FieldBinlog
|
||||
statsLog *datapb.FieldBinlog
|
||||
bm25StatsLog map[FieldID]*datapb.FieldBinlog
|
||||
}
|
||||
|
||||
var _ BinlogRecordWriter = (*CompositeBinlogRecordWriter)(nil)
|
||||
|
||||
func (c *CompositeBinlogRecordWriter) Write(r Record) error {
|
||||
if err := c.initWriters(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
tsArray := r.Column(common.TimeStampField).(*array.Int64)
|
||||
rows := r.Len()
|
||||
for i := 0; i < rows; i++ {
|
||||
ts := typeutil.Timestamp(tsArray.Value(i))
|
||||
if ts < c.tsFrom {
|
||||
c.tsFrom = ts
|
||||
}
|
||||
if ts > c.tsTo {
|
||||
c.tsTo = ts
|
||||
}
|
||||
|
||||
switch schemapb.DataType(c.pkstats.PkType) {
|
||||
case schemapb.DataType_Int64:
|
||||
pkArray := r.Column(c.pkstats.FieldID).(*array.Int64)
|
||||
pk := &Int64PrimaryKey{
|
||||
Value: pkArray.Value(i),
|
||||
}
|
||||
c.pkstats.Update(pk)
|
||||
case schemapb.DataType_VarChar:
|
||||
pkArray := r.Column(c.pkstats.FieldID).(*array.String)
|
||||
pk := &VarCharPrimaryKey{
|
||||
Value: pkArray.Value(i),
|
||||
}
|
||||
c.pkstats.Update(pk)
|
||||
default:
|
||||
panic("invalid data type")
|
||||
}
|
||||
|
||||
for fieldID, stats := range c.bm25Stats {
|
||||
field, ok := r.Column(fieldID).(*array.Binary)
|
||||
if !ok {
|
||||
return fmt.Errorf("bm25 field value not found")
|
||||
}
|
||||
stats.AppendBytes(field.Value(i))
|
||||
}
|
||||
}
|
||||
|
||||
if err := c.rw.Write(r); err != nil {
|
||||
return err
|
||||
}
|
||||
c.rowNum += int64(rows)
|
||||
|
||||
// flush if size exceeds chunk size
|
||||
if c.rw.GetWrittenUncompressed() >= c.chunkSize {
|
||||
return c.flushChunk()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *CompositeBinlogRecordWriter) initWriters() error {
|
||||
if c.rw == nil {
|
||||
c.fieldWriters = NewBinlogStreamWriters(c.collectionID, c.partitionID, c.segmentID, c.schema.Fields)
|
||||
rws := make(map[FieldID]RecordWriter, len(c.fieldWriters))
|
||||
for fid, w := range c.fieldWriters {
|
||||
rw, err := w.GetRecordWriter()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
rws[fid] = rw
|
||||
}
|
||||
c.rw = NewCompositeRecordWriter(rws)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *CompositeBinlogRecordWriter) resetWriters() {
|
||||
c.fieldWriters = nil
|
||||
c.rw = nil
|
||||
c.tsFrom = math.MaxUint64
|
||||
c.tsTo = 0
|
||||
}
|
||||
|
||||
func (c *CompositeBinlogRecordWriter) Close() error {
|
||||
if err := c.writeStats(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := c.writeBm25Stats(); err != nil {
|
||||
return err
|
||||
}
|
||||
if c.rw != nil {
|
||||
// if rw is not nil, it means there is data to be flushed
|
||||
if err := c.flushChunk(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *CompositeBinlogRecordWriter) GetWrittenUncompressed() uint64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
func (c *CompositeBinlogRecordWriter) flushChunk() error {
|
||||
if c.fieldWriters == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
id, _, err := c.allocator.Alloc(uint32(len(c.fieldWriters)))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
blobs := make(map[FieldID]*Blob, len(c.fieldWriters))
|
||||
for fid, w := range c.fieldWriters {
|
||||
b, err := w.Finalize()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// assign blob key
|
||||
b.Key = metautil.BuildInsertLogPath(c.rootPath, c.collectionID, c.partitionID, c.segmentID, fid, id)
|
||||
blobs[fid] = b
|
||||
id++
|
||||
}
|
||||
if err := c.BlobsWriter(lo.Values(blobs)); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// attach binlogs
|
||||
if c.fieldBinlogs == nil {
|
||||
c.fieldBinlogs = make(map[FieldID]*datapb.FieldBinlog, len(c.fieldWriters))
|
||||
for fid := range c.fieldWriters {
|
||||
c.fieldBinlogs[fid] = &datapb.FieldBinlog{
|
||||
FieldID: fid,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for fid, b := range blobs {
|
||||
c.fieldBinlogs[fid].Binlogs = append(c.fieldBinlogs[fid].Binlogs, &datapb.Binlog{
|
||||
LogSize: int64(len(b.Value)),
|
||||
MemorySize: b.MemorySize,
|
||||
LogPath: b.Key,
|
||||
EntriesNum: b.RowNum,
|
||||
TimestampFrom: c.tsFrom,
|
||||
TimestampTo: c.tsTo,
|
||||
})
|
||||
}
|
||||
|
||||
// reset writers
|
||||
c.resetWriters()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *CompositeBinlogRecordWriter) writeStats() error {
|
||||
if c.pkstats == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
id, err := c.allocator.AllocOne()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
codec := NewInsertCodecWithSchema(&etcdpb.CollectionMeta{
|
||||
ID: c.collectionID,
|
||||
Schema: c.schema,
|
||||
})
|
||||
sblob, err := codec.SerializePkStats(c.pkstats, c.rowNum)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
sblob.Key = metautil.BuildStatsLogPath(c.rootPath,
|
||||
c.collectionID, c.partitionID, c.segmentID, c.pkstats.FieldID, id)
|
||||
|
||||
if err := c.BlobsWriter([]*Blob{sblob}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
c.statsLog = &datapb.FieldBinlog{
|
||||
FieldID: c.pkstats.FieldID,
|
||||
Binlogs: []*datapb.Binlog{
|
||||
{
|
||||
LogSize: int64(len(sblob.GetValue())),
|
||||
MemorySize: int64(len(sblob.GetValue())),
|
||||
LogPath: sblob.Key,
|
||||
EntriesNum: c.rowNum,
|
||||
},
|
||||
},
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *CompositeBinlogRecordWriter) writeBm25Stats() error {
|
||||
if len(c.bm25Stats) == 0 {
|
||||
return nil
|
||||
}
|
||||
id, _, err := c.allocator.Alloc(uint32(len(c.bm25Stats)))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if c.bm25StatsLog == nil {
|
||||
c.bm25StatsLog = make(map[FieldID]*datapb.FieldBinlog)
|
||||
}
|
||||
for fid, stats := range c.bm25Stats {
|
||||
bytes, err := stats.Serialize()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key := metautil.BuildBm25LogPath(c.rootPath,
|
||||
c.collectionID, c.partitionID, c.segmentID, fid, id)
|
||||
blob := &Blob{
|
||||
Key: key,
|
||||
Value: bytes,
|
||||
RowNum: stats.NumRow(),
|
||||
MemorySize: int64(len(bytes)),
|
||||
}
|
||||
if err := c.BlobsWriter([]*Blob{blob}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
fieldLog := &datapb.FieldBinlog{
|
||||
FieldID: fid,
|
||||
Binlogs: []*datapb.Binlog{
|
||||
{
|
||||
LogSize: int64(len(blob.GetValue())),
|
||||
MemorySize: int64(len(blob.GetValue())),
|
||||
LogPath: key,
|
||||
EntriesNum: c.rowNum,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
c.bm25StatsLog[fid] = fieldLog
|
||||
id++
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *CompositeBinlogRecordWriter) GetLogs() (
|
||||
fieldBinlogs map[FieldID]*datapb.FieldBinlog,
|
||||
statsLog *datapb.FieldBinlog,
|
||||
bm25StatsLog map[FieldID]*datapb.FieldBinlog,
|
||||
) {
|
||||
return c.fieldBinlogs, c.statsLog, c.bm25StatsLog
|
||||
}
|
||||
|
||||
func (c *CompositeBinlogRecordWriter) GetRowNum() int64 {
|
||||
return c.rowNum
|
||||
}
|
||||
|
||||
func newCompositeBinlogRecordWriter(collectionID, partitionID, segmentID UniqueID, schema *schemapb.CollectionSchema,
|
||||
blobsWriter ChunkedBlobsWriter, allocator allocator.Interface, chunkSize uint64, rootPath string, maxRowNum int64,
|
||||
) (*CompositeBinlogRecordWriter, error) {
|
||||
pkField, err := typeutil.GetPrimaryFieldSchema(schema)
|
||||
if err != nil {
|
||||
log.Warn("failed to get pk field from schema")
|
||||
return nil, err
|
||||
}
|
||||
stats, err := NewPrimaryKeyStats(pkField.GetFieldID(), int64(pkField.GetDataType()), maxRowNum)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
bm25FieldIDs := lo.FilterMap(schema.GetFunctions(), func(function *schemapb.FunctionSchema, _ int) (int64, bool) {
|
||||
if function.GetType() == schemapb.FunctionType_BM25 {
|
||||
return function.GetOutputFieldIds()[0], true
|
||||
}
|
||||
return 0, false
|
||||
})
|
||||
bm25Stats := make(map[int64]*BM25Stats, len(bm25FieldIDs))
|
||||
for _, fid := range bm25FieldIDs {
|
||||
bm25Stats[fid] = NewBM25Stats()
|
||||
}
|
||||
|
||||
return &CompositeBinlogRecordWriter{
|
||||
collectionID: collectionID,
|
||||
partitionID: partitionID,
|
||||
segmentID: segmentID,
|
||||
schema: schema,
|
||||
BlobsWriter: blobsWriter,
|
||||
allocator: allocator,
|
||||
chunkSize: chunkSize,
|
||||
rootPath: rootPath,
|
||||
maxRowNum: maxRowNum,
|
||||
pkstats: stats,
|
||||
bm25Stats: bm25Stats,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func NewChunkedBinlogSerializeWriter(collectionID, partitionID, segmentID UniqueID, schema *schemapb.CollectionSchema,
|
||||
blobsWriter ChunkedBlobsWriter, allocator allocator.Interface, chunkSize uint64, rootPath string, maxRowNum int64, batchSize int,
|
||||
) (*SerializeWriter[*Value], error) {
|
||||
rw, err := newCompositeBinlogRecordWriter(collectionID, partitionID, segmentID, schema, blobsWriter, allocator, chunkSize, rootPath, maxRowNum)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return NewSerializeRecordWriter[*Value](rw, func(v []*Value) (Record, error) {
|
||||
return ValueSerializer(v, schema.Fields)
|
||||
}, batchSize), nil
|
||||
}
|
||||
|
||||
func NewBinlogSerializeWriter(schema *schemapb.CollectionSchema, partitionID, segmentID UniqueID,
|
||||
eventWriters map[FieldID]*BinlogStreamWriter, batchSize int,
|
||||
) (*SerializeWriter[*Value], error) {
|
||||
@ -623,17 +966,17 @@ func (crr *simpleArrowRecordReader) iterateNextBatch() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (crr *simpleArrowRecordReader) Next() error {
|
||||
func (crr *simpleArrowRecordReader) Next() (Record, error) {
|
||||
if crr.rr == nil {
|
||||
if len(crr.blobs) == 0 {
|
||||
return io.EOF
|
||||
return nil, io.EOF
|
||||
}
|
||||
crr.blobPos = -1
|
||||
crr.r = simpleArrowRecord{
|
||||
field2Col: make(map[FieldID]int),
|
||||
}
|
||||
if err := crr.iterateNextBatch(); err != nil {
|
||||
return err
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
@ -651,17 +994,13 @@ func (crr *simpleArrowRecordReader) Next() error {
|
||||
|
||||
if ok := composeRecord(); !ok {
|
||||
if err := crr.iterateNextBatch(); err != nil {
|
||||
return err
|
||||
return nil, err
|
||||
}
|
||||
if ok := composeRecord(); !ok {
|
||||
return io.EOF
|
||||
return nil, io.EOF
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (crr *simpleArrowRecordReader) Record() Record {
|
||||
return &crr.r
|
||||
return &crr.r, nil
|
||||
}
|
||||
|
||||
func (crr *simpleArrowRecordReader) Close() error {
|
||||
|
||||
@ -33,24 +33,20 @@ type packedRecordReader struct {
|
||||
|
||||
bufferSize int64
|
||||
schema *schemapb.CollectionSchema
|
||||
r *simpleArrowRecord
|
||||
field2Col map[FieldID]int
|
||||
}
|
||||
|
||||
func (pr *packedRecordReader) Next() error {
|
||||
var _ RecordReader = (*packedRecordReader)(nil)
|
||||
|
||||
func (pr *packedRecordReader) Next() (Record, error) {
|
||||
if pr.reader == nil {
|
||||
return io.EOF
|
||||
return nil, io.EOF
|
||||
}
|
||||
rec, err := pr.reader.ReadNext()
|
||||
if err != nil || rec == nil {
|
||||
return io.EOF
|
||||
return nil, io.EOF
|
||||
}
|
||||
pr.r = NewSimpleArrowRecord(rec, pr.field2Col)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (pr *packedRecordReader) Record() Record {
|
||||
return pr.r
|
||||
return NewSimpleArrowRecord(rec, pr.field2Col), nil
|
||||
}
|
||||
|
||||
func (pr *packedRecordReader) Close() error {
|
||||
@ -60,7 +56,7 @@ func (pr *packedRecordReader) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func NewPackedRecordReader(paths []string, schema *schemapb.CollectionSchema, bufferSize int64,
|
||||
func newPackedRecordReader(paths []string, schema *schemapb.CollectionSchema, bufferSize int64,
|
||||
) (*packedRecordReader, error) {
|
||||
arrowSchema, err := ConvertToArrowSchema(schema.Fields)
|
||||
if err != nil {
|
||||
@ -85,7 +81,7 @@ func NewPackedRecordReader(paths []string, schema *schemapb.CollectionSchema, bu
|
||||
func NewPackedDeserializeReader(paths []string, schema *schemapb.CollectionSchema,
|
||||
bufferSize int64, pkFieldID FieldID,
|
||||
) (*DeserializeReader[*Value], error) {
|
||||
reader, err := NewPackedRecordReader(paths, schema, bufferSize)
|
||||
reader, err := newPackedRecordReader(paths, schema, bufferSize)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@ -27,10 +27,11 @@ import (
|
||||
"github.com/apache/arrow/go/v12/arrow/memory"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/util/typeutil"
|
||||
)
|
||||
|
||||
func Sort(schema *schemapb.CollectionSchema, rr []RecordReader,
|
||||
pkField FieldID, rw RecordWriter, predicate func(r Record, ri, i int) bool,
|
||||
rw RecordWriter, predicate func(r Record, ri, i int) bool,
|
||||
) (int, error) {
|
||||
records := make([]Record, 0)
|
||||
|
||||
@ -48,9 +49,8 @@ func Sort(schema *schemapb.CollectionSchema, rr []RecordReader,
|
||||
|
||||
for _, r := range rr {
|
||||
for {
|
||||
err := r.Next()
|
||||
rec, err := r.Next()
|
||||
if err == nil {
|
||||
rec := r.Record()
|
||||
rec.Retain()
|
||||
ri := len(records)
|
||||
records = append(records, rec)
|
||||
@ -71,17 +71,23 @@ func Sort(schema *schemapb.CollectionSchema, rr []RecordReader,
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
switch records[0].Column(pkField).(type) {
|
||||
pkField, err := typeutil.GetPrimaryFieldSchema(schema)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
pkFieldId := pkField.FieldID
|
||||
|
||||
switch records[0].Column(pkFieldId).(type) {
|
||||
case *array.Int64:
|
||||
sort.Slice(indices, func(i, j int) bool {
|
||||
pki := records[indices[i].ri].Column(pkField).(*array.Int64).Value(indices[i].i)
|
||||
pkj := records[indices[j].ri].Column(pkField).(*array.Int64).Value(indices[j].i)
|
||||
pki := records[indices[i].ri].Column(pkFieldId).(*array.Int64).Value(indices[i].i)
|
||||
pkj := records[indices[j].ri].Column(pkFieldId).(*array.Int64).Value(indices[j].i)
|
||||
return pki < pkj
|
||||
})
|
||||
case *array.String:
|
||||
sort.Slice(indices, func(i, j int) bool {
|
||||
pki := records[indices[i].ri].Column(pkField).(*array.String).Value(indices[i].i)
|
||||
pkj := records[indices[j].ri].Column(pkField).(*array.String).Value(indices[j].i)
|
||||
pki := records[indices[i].ri].Column(pkFieldId).(*array.String).Value(indices[i].i)
|
||||
pkj := records[indices[j].ri].Column(pkFieldId).(*array.String).Value(indices[j].i)
|
||||
return pki < pkj
|
||||
})
|
||||
}
|
||||
@ -192,47 +198,53 @@ func NewPriorityQueue[T any](less func(x, y *T) bool) *PriorityQueue[T] {
|
||||
}
|
||||
|
||||
func MergeSort(schema *schemapb.CollectionSchema, rr []RecordReader,
|
||||
pkField FieldID, rw RecordWriter, predicate func(r Record, ri, i int) bool,
|
||||
rw RecordWriter, predicate func(r Record, ri, i int) bool,
|
||||
) (numRows int, err error) {
|
||||
type index struct {
|
||||
ri int
|
||||
i int
|
||||
}
|
||||
|
||||
advanceRecord := func(r RecordReader) (Record, error) {
|
||||
err := r.Next()
|
||||
recs := make([]Record, len(rr))
|
||||
advanceRecord := func(i int) error {
|
||||
rec, err := rr[i].Next()
|
||||
recs[i] = rec // assign nil if err
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return err
|
||||
}
|
||||
return r.Record(), nil
|
||||
return nil
|
||||
}
|
||||
|
||||
recs := make([]Record, len(rr))
|
||||
for i, r := range rr {
|
||||
rec, err := advanceRecord(r)
|
||||
for i := range rr {
|
||||
err := advanceRecord(i)
|
||||
if err == io.EOF {
|
||||
recs[i] = nil
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
recs[i] = rec
|
||||
}
|
||||
|
||||
pkField, err := typeutil.GetPrimaryFieldSchema(schema)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
pkFieldId := pkField.FieldID
|
||||
|
||||
var pq *PriorityQueue[index]
|
||||
switch recs[0].Column(pkField).(type) {
|
||||
switch recs[0].Column(pkFieldId).(type) {
|
||||
case *array.Int64:
|
||||
pq = NewPriorityQueue[index](func(x, y *index) bool {
|
||||
return rr[x.ri].Record().Column(pkField).(*array.Int64).Value(x.i) < rr[y.ri].Record().Column(pkField).(*array.Int64).Value(y.i)
|
||||
return recs[x.ri].Column(pkFieldId).(*array.Int64).Value(x.i) < recs[y.ri].Column(pkFieldId).(*array.Int64).Value(y.i)
|
||||
})
|
||||
case *array.String:
|
||||
pq = NewPriorityQueue[index](func(x, y *index) bool {
|
||||
return rr[x.ri].Record().Column(pkField).(*array.String).Value(x.i) < rr[y.ri].Record().Column(pkField).(*array.String).Value(y.i)
|
||||
return recs[x.ri].Column(pkFieldId).(*array.String).Value(x.i) < recs[y.ri].Column(pkFieldId).(*array.String).Value(y.i)
|
||||
})
|
||||
}
|
||||
|
||||
enqueueAll := func(ri int, r Record) {
|
||||
enqueueAll := func(ri int) {
|
||||
r := recs[ri]
|
||||
for j := 0; j < r.Len(); j++ {
|
||||
if predicate(r, ri, j) {
|
||||
pq.Enqueue(&index{
|
||||
@ -246,7 +258,7 @@ func MergeSort(schema *schemapb.CollectionSchema, rr []RecordReader,
|
||||
|
||||
for i, v := range recs {
|
||||
if v != nil {
|
||||
enqueueAll(i, v)
|
||||
enqueueAll(i)
|
||||
}
|
||||
}
|
||||
|
||||
@ -294,7 +306,7 @@ func MergeSort(schema *schemapb.CollectionSchema, rr []RecordReader,
|
||||
for c, builder := range builders {
|
||||
fid := schema.Fields[c].FieldID
|
||||
defaultValue := schema.Fields[c].GetDefaultValue()
|
||||
AppendValueAt(builder, rr[idx.ri].Record().Column(fid), idx.i, defaultValue)
|
||||
AppendValueAt(builder, recs[idx.ri].Column(fid), idx.i, defaultValue)
|
||||
}
|
||||
if (rc+1)%batchSize == 0 {
|
||||
writeRecord(int64(batchSize))
|
||||
@ -304,15 +316,15 @@ func MergeSort(schema *schemapb.CollectionSchema, rr []RecordReader,
|
||||
}
|
||||
|
||||
// If poped idx reaches end of segment, invalidate cache and advance to next segment
|
||||
if idx.i == rr[idx.ri].Record().Len()-1 {
|
||||
rec, err := advanceRecord(rr[idx.ri])
|
||||
if idx.i == recs[idx.ri].Len()-1 {
|
||||
err := advanceRecord(idx.ri)
|
||||
if err == io.EOF {
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
enqueueAll(idx.ri, rec)
|
||||
enqueueAll(idx.ri)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -29,11 +29,11 @@ func TestSort(t *testing.T) {
|
||||
getReaders := func() []RecordReader {
|
||||
blobs, err := generateTestDataWithSeed(10, 3)
|
||||
assert.NoError(t, err)
|
||||
reader10, err := NewCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
reader10, err := newCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
assert.NoError(t, err)
|
||||
blobs, err = generateTestDataWithSeed(20, 3)
|
||||
assert.NoError(t, err)
|
||||
reader20, err := NewCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
reader20, err := newCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
assert.NoError(t, err)
|
||||
rr := []RecordReader{reader20, reader10}
|
||||
return rr
|
||||
@ -55,7 +55,7 @@ func TestSort(t *testing.T) {
|
||||
}
|
||||
|
||||
t.Run("sort", func(t *testing.T) {
|
||||
gotNumRows, err := Sort(generateTestSchema(), getReaders(), common.RowIDField, rw, func(r Record, ri, i int) bool {
|
||||
gotNumRows, err := Sort(generateTestSchema(), getReaders(), rw, func(r Record, ri, i int) bool {
|
||||
return true
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
@ -65,7 +65,7 @@ func TestSort(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("sort with predicate", func(t *testing.T) {
|
||||
gotNumRows, err := Sort(generateTestSchema(), getReaders(), common.RowIDField, rw, func(r Record, ri, i int) bool {
|
||||
gotNumRows, err := Sort(generateTestSchema(), getReaders(), rw, func(r Record, ri, i int) bool {
|
||||
pk := r.Column(common.RowIDField).(*array.Int64).Value(i)
|
||||
return pk >= 20
|
||||
})
|
||||
@ -80,11 +80,11 @@ func TestMergeSort(t *testing.T) {
|
||||
getReaders := func() []RecordReader {
|
||||
blobs, err := generateTestDataWithSeed(10, 3)
|
||||
assert.NoError(t, err)
|
||||
reader10, err := NewCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
reader10, err := newCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
assert.NoError(t, err)
|
||||
blobs, err = generateTestDataWithSeed(20, 3)
|
||||
assert.NoError(t, err)
|
||||
reader20, err := NewCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
reader20, err := newCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
assert.NoError(t, err)
|
||||
rr := []RecordReader{reader20, reader10}
|
||||
return rr
|
||||
@ -106,7 +106,7 @@ func TestMergeSort(t *testing.T) {
|
||||
}
|
||||
|
||||
t.Run("merge sort", func(t *testing.T) {
|
||||
gotNumRows, err := MergeSort(generateTestSchema(), getReaders(), common.RowIDField, rw, func(r Record, ri, i int) bool {
|
||||
gotNumRows, err := MergeSort(generateTestSchema(), getReaders(), rw, func(r Record, ri, i int) bool {
|
||||
return true
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
@ -116,7 +116,7 @@ func TestMergeSort(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("merge sort with predicate", func(t *testing.T) {
|
||||
gotNumRows, err := MergeSort(generateTestSchema(), getReaders(), common.RowIDField, rw, func(r Record, ri, i int) bool {
|
||||
gotNumRows, err := MergeSort(generateTestSchema(), getReaders(), rw, func(r Record, ri, i int) bool {
|
||||
pk := r.Column(common.RowIDField).(*array.Int64).Value(i)
|
||||
return pk >= 20
|
||||
})
|
||||
@ -132,11 +132,11 @@ func BenchmarkSort(b *testing.B) {
|
||||
batch := 500000
|
||||
blobs, err := generateTestDataWithSeed(batch, batch)
|
||||
assert.NoError(b, err)
|
||||
reader10, err := NewCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
reader10, err := newCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
assert.NoError(b, err)
|
||||
blobs, err = generateTestDataWithSeed(batch*2+1, batch)
|
||||
assert.NoError(b, err)
|
||||
reader20, err := NewCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
reader20, err := newCompositeBinlogRecordReader(generateTestSchema(), MakeBlobsReader(blobs))
|
||||
assert.NoError(b, err)
|
||||
rr := []RecordReader{reader20, reader10}
|
||||
|
||||
@ -154,7 +154,7 @@ func BenchmarkSort(b *testing.B) {
|
||||
|
||||
b.Run("sort", func(b *testing.B) {
|
||||
for i := 0; i < b.N; i++ {
|
||||
Sort(generateTestSchema(), rr, common.RowIDField, rw, func(r Record, ri, i int) bool {
|
||||
Sort(generateTestSchema(), rr, rw, func(r Record, ri, i int) bool {
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
@ -594,6 +594,7 @@ message CompactionSegmentBinlogs {
|
||||
int64 collectionID = 7;
|
||||
int64 partitionID = 8;
|
||||
bool is_sorted = 9;
|
||||
int64 storage_version = 10;
|
||||
}
|
||||
|
||||
message CompactionPlan {
|
||||
@ -631,6 +632,7 @@ message CompactionSegment {
|
||||
string channel = 7;
|
||||
bool is_sorted = 8;
|
||||
repeated FieldBinlog bm25logs = 9;
|
||||
int64 storage_version = 10;
|
||||
}
|
||||
|
||||
message CompactionPlanResult {
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Loading…
x
Reference in New Issue
Block a user