wehub-resource-sync
498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
2684 lines
96 KiB
Go
2684 lines
96 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 segments
|
|
|
|
/*
|
|
#cgo pkg-config: milvus_core
|
|
|
|
#include "segcore/load_index_c.h"
|
|
*/
|
|
import "C"
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"math"
|
|
"path"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/apache/arrow/go/v17/arrow/array"
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/samber/lo"
|
|
"go.opentelemetry.io/otel"
|
|
"go.uber.org/atomic"
|
|
"golang.org/x/sync/errgroup"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
|
"github.com/milvus-io/milvus/internal/querynodev2/pkoracle"
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
"github.com/milvus-io/milvus/internal/storagecommon"
|
|
"github.com/milvus-io/milvus/internal/storagev2/packed"
|
|
"github.com/milvus-io/milvus/internal/util/indexparamcheck"
|
|
"github.com/milvus-io/milvus/internal/util/vecindexmgr"
|
|
"github.com/milvus-io/milvus/pkg/v3/common"
|
|
"github.com/milvus-io/milvus/pkg/v3/metrics"
|
|
"github.com/milvus-io/milvus/pkg/v3/mlog"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/hardware"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/indexparams"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/logutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/metautil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/metric"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/syncutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
|
|
)
|
|
|
|
const (
|
|
UsedDiskMemoryRatio = 4
|
|
UsedDiskMemoryRatioAisaq = 64
|
|
)
|
|
|
|
var errRetryTimerNotified = errors.New("retry timer notified")
|
|
|
|
type Loader interface {
|
|
// Load loads binlogs, and spawn segments,
|
|
// NOTE: make sure the ref count of the corresponding collection will never go down to 0 during this
|
|
Load(ctx context.Context, collectionID int64, segmentType SegmentType, version int64, segments ...*querypb.SegmentLoadInfo) ([]Segment, error)
|
|
|
|
// LoadDeltaLogs load deltalog and write delta data into provided segment.
|
|
// it also executes resource protection logic in case of OOM.
|
|
LoadDeltaLogs(ctx context.Context, segment Segment, loadInfo *querypb.SegmentLoadInfo) error
|
|
|
|
// LoadBloomFilterSet loads needed statslog for RemoteSegment.
|
|
LoadBloomFilterSet(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) ([]*pkoracle.BloomFilterSet, error)
|
|
|
|
// GetChunkManager returns the chunk manager for remote storage access.
|
|
GetChunkManager() storage.ChunkManager
|
|
|
|
// LoadIndex append index for segment and remove vector binlogs.
|
|
LoadIndex(ctx context.Context,
|
|
segment Segment,
|
|
info *querypb.SegmentLoadInfo,
|
|
version int64) error
|
|
|
|
// ReopenSegments update segment data according to new load info.
|
|
ReopenSegments(ctx context.Context,
|
|
loadInfos []*querypb.SegmentLoadInfo,
|
|
) error
|
|
}
|
|
|
|
type ResourceEstimate struct {
|
|
MaxMemoryCost uint64
|
|
MaxDiskCost uint64
|
|
FinalMemoryCost uint64
|
|
FinalDiskCost uint64
|
|
HasRawData bool
|
|
}
|
|
|
|
func GetResourceEstimate(estimate *C.LoadResourceRequest) ResourceEstimate {
|
|
return ResourceEstimate{
|
|
MaxMemoryCost: uint64(estimate.max_memory_cost),
|
|
MaxDiskCost: uint64(estimate.max_disk_cost),
|
|
FinalMemoryCost: uint64(estimate.final_memory_cost),
|
|
FinalDiskCost: uint64(estimate.final_disk_cost),
|
|
HasRawData: bool(estimate.has_raw_data),
|
|
}
|
|
}
|
|
|
|
type requestResourceResult struct {
|
|
Resource LoadResource
|
|
LogicalResource LoadResource
|
|
CommittedResource LoadResource
|
|
ConcurrencyLevel int
|
|
}
|
|
|
|
type LoadResource struct {
|
|
MemorySize uint64
|
|
DiskSize uint64
|
|
}
|
|
|
|
func (r *LoadResource) Add(resource LoadResource) {
|
|
r.MemorySize += resource.MemorySize
|
|
r.DiskSize += resource.DiskSize
|
|
}
|
|
|
|
func (r *LoadResource) Sub(resource LoadResource) {
|
|
r.MemorySize -= resource.MemorySize
|
|
r.DiskSize -= resource.DiskSize
|
|
}
|
|
|
|
func (r *LoadResource) IsZero() bool {
|
|
return r.MemorySize == 0 && r.DiskSize == 0
|
|
}
|
|
|
|
type resourceEstimateFactor struct {
|
|
memoryUsageFactor float64
|
|
memoryIndexUsageFactor float64
|
|
EnableInterminSegmentIndex bool
|
|
tempSegmentIndexFactor float64
|
|
deltaDataExpansionFactor float64
|
|
jsonKeyStatsExpansionFactor float64
|
|
textIndexExpansionFactor float64
|
|
TieredEvictionEnabled bool
|
|
TieredEvictableMemoryCacheRatio float64
|
|
TieredEvictableDiskCacheRatio float64
|
|
// externalRawDataFactor is the peak-memory safety factor for external
|
|
// segments. External tables always download, decompress and deserialize
|
|
// row groups into Arrow buffers regardless of mmap / TieredEviction
|
|
// settings, so peak transient memory = rawDataSize * factor. Defaults
|
|
// to 2.0 via paramtable queryNode.externalCollection.rawDataFactor.
|
|
externalRawDataFactor float64
|
|
}
|
|
|
|
func NewLoader(
|
|
ctx context.Context,
|
|
manager *Manager,
|
|
cm storage.ChunkManager,
|
|
) *segmentLoader {
|
|
duf := NewDiskUsageFetcher(ctx)
|
|
go duf.Start()
|
|
|
|
loader := &segmentLoader{
|
|
manager: manager,
|
|
cm: cm,
|
|
loadingSegments: typeutil.NewConcurrentMap[int64, *loadResult](),
|
|
committedResourceNotifier: syncutil.NewVersionedNotifier(),
|
|
duf: duf,
|
|
}
|
|
|
|
return loader
|
|
}
|
|
|
|
type loadStatus = int32
|
|
|
|
const (
|
|
loading loadStatus = iota + 1
|
|
success
|
|
failure
|
|
)
|
|
|
|
type loadResult struct {
|
|
status *atomic.Int32
|
|
cond *sync.Cond
|
|
}
|
|
|
|
func newLoadResult() *loadResult {
|
|
return &loadResult{
|
|
status: atomic.NewInt32(loading),
|
|
cond: sync.NewCond(&sync.Mutex{}),
|
|
}
|
|
}
|
|
|
|
func (r *loadResult) SetResult(status loadStatus) {
|
|
r.status.CompareAndSwap(loading, status)
|
|
r.cond.Broadcast()
|
|
}
|
|
|
|
// segmentLoader is only responsible for loading the field data from binlog
|
|
type segmentLoader struct {
|
|
manager *Manager
|
|
cm storage.ChunkManager
|
|
|
|
// The channel will be closed as the segment loaded
|
|
loadingSegments *typeutil.ConcurrentMap[int64, *loadResult]
|
|
|
|
mut sync.Mutex // guards committedResource
|
|
committedResource LoadResource
|
|
committedLogicalResource LoadResource
|
|
committedResourceNotifier *syncutil.VersionedNotifier
|
|
|
|
duf *diskUsageFetcher
|
|
}
|
|
|
|
var _ Loader = (*segmentLoader)(nil)
|
|
|
|
func (loader *segmentLoader) Load(ctx context.Context,
|
|
collectionID int64,
|
|
segmentType SegmentType,
|
|
version int64,
|
|
segments ...*querypb.SegmentLoadInfo,
|
|
) ([]Segment, error) {
|
|
if len(segments) == 0 {
|
|
mlog.Info(context.TODO(), "no segment to load")
|
|
return nil, nil
|
|
}
|
|
|
|
collection := loader.manager.Collection.Get(collectionID)
|
|
if collection == nil {
|
|
err := merr.WrapErrCollectionNotFound(collectionID)
|
|
mlog.Warn(context.TODO(), "failed to get collection", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
for _, segment := range segments {
|
|
configureUseTakeForOutput(segment, collection.Schema())
|
|
}
|
|
// Filter out loaded & loading segments
|
|
infos := loader.prepare(ctx, segmentType, segments...)
|
|
defer loader.unregister(infos...)
|
|
|
|
// continue to wait other task done
|
|
mlog.Info(context.TODO(), "start loading...", mlog.Int("segmentNum", len(segments)), mlog.Int("afterFilter", len(infos)))
|
|
|
|
var err error
|
|
var requestResourceResult requestResourceResult
|
|
|
|
// Check memory & storage limit
|
|
// no need to check resource for lazy load here
|
|
requestResourceResult, err = loader.requestResource(ctx, infos...)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "request resource failed", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
defer loader.freeRequestResource(requestResourceResult)
|
|
|
|
newSegments := typeutil.NewConcurrentMap[int64, Segment]()
|
|
loaded := typeutil.NewConcurrentMap[int64, Segment]()
|
|
defer func() {
|
|
newSegments.Range(func(segmentID int64, s Segment) bool {
|
|
mlog.Warn(context.TODO(), "release new segment created due to load failure",
|
|
mlog.Int64("segmentID", segmentID),
|
|
mlog.Err(err),
|
|
)
|
|
s.Release(context.Background())
|
|
return true
|
|
})
|
|
}()
|
|
|
|
for _, info := range infos {
|
|
loadInfo := info
|
|
|
|
for _, indexInfo := range loadInfo.IndexInfos {
|
|
indexParams := funcutil.KeyValuePair2Map(indexInfo.IndexParams)
|
|
|
|
// some build params also exist in indexParams, which are useless during loading process
|
|
if vecindexmgr.GetVecIndexMgrInstance().IsDiskANN(indexParams["index_type"]) {
|
|
if err := indexparams.SetDiskIndexLoadParams(paramtable.Get(), indexParams, indexInfo.GetNumRows()); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
// set whether enable offset cache for bitmap index
|
|
if indexParams["index_type"] == indexparamcheck.IndexBitmap {
|
|
indexparams.SetBitmapIndexLoadParams(paramtable.Get(), indexParams)
|
|
}
|
|
|
|
if err := indexparams.AppendPrepareLoadParams(paramtable.Get(), indexParams); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
indexInfo.IndexParams = funcutil.Map2KeyValuePair(indexParams)
|
|
}
|
|
|
|
segment, err := NewSegment(
|
|
ctx,
|
|
collection,
|
|
loader.manager.Segment,
|
|
segmentType,
|
|
version,
|
|
loadInfo,
|
|
)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "load segment failed when create new segment",
|
|
mlog.Int64("partitionID", loadInfo.GetPartitionID()),
|
|
mlog.Int64("segmentID", loadInfo.GetSegmentID()),
|
|
mlog.Err(err),
|
|
)
|
|
return nil, err
|
|
}
|
|
|
|
newSegments.Insert(loadInfo.GetSegmentID(), segment)
|
|
}
|
|
|
|
loadSegmentFunc := func(idx int) (err error) {
|
|
loadInfo := infos[idx]
|
|
partitionID := loadInfo.PartitionID
|
|
segmentID := loadInfo.SegmentID
|
|
segment, _ := newSegments.Get(segmentID)
|
|
|
|
logger := mlog.With(mlog.Int64("partitionID", partitionID),
|
|
mlog.Int64("segmentID", segmentID),
|
|
mlog.String("segmentType", loadInfo.GetLevel().String()))
|
|
metrics.QueryNodeLoadSegmentConcurrency.WithLabelValues(paramtable.GetStringNodeID(), "LoadSegment").Inc()
|
|
defer func() {
|
|
metrics.QueryNodeLoadSegmentConcurrency.WithLabelValues(paramtable.GetStringNodeID(), "LoadSegment").Dec()
|
|
if err != nil {
|
|
logger.Warn(ctx, "load segment failed when load data into memory", mlog.Err(err))
|
|
}
|
|
logger.Info(ctx, "load segment done")
|
|
}()
|
|
tr := timerecord.NewTimeRecorder("loadDurationPerSegment")
|
|
logger.Info(ctx, "load segment...")
|
|
|
|
// L0 segment has no index or data to be load.
|
|
if loadInfo.GetLevel() != datapb.SegmentLevel_L0 {
|
|
// lazy load segment do not load segment at first time.
|
|
if err = loader.LoadSegment(ctx, segment, loadInfo); err != nil {
|
|
return merr.Wrap(err, "At LoadSegment")
|
|
}
|
|
}
|
|
if err = loader.loadDeltalogs(ctx, segment, loadInfo); err != nil {
|
|
return merr.Wrap(err, "At LoadDeltaLogs")
|
|
}
|
|
|
|
schema := collection.Schema()
|
|
isExternalCollection := typeutil.IsExternalCollection(schema)
|
|
isMilvusTableRealPK := typeutil.NewStorageColumnResolver(schema).IsMilvusTable() &&
|
|
HasExternalPrimaryKey(schema)
|
|
if !segment.PkCandidateExist() {
|
|
mlog.Debug(context.TODO(), "loading PK candidate for segment", mlog.Int64("segmentID", segment.ID()))
|
|
if isExternalCollection {
|
|
var candidate pkoracle.Candidate
|
|
if isMilvusTableRealPK {
|
|
bfs, err := loader.loadSingleBloomFilterSet(ctx, loadInfo.GetCollectionID(), loadInfo, segment.Type())
|
|
if err != nil {
|
|
return merr.Wrap(err, "At LoadBloomFilter")
|
|
}
|
|
if bfs.PkCandidateExist() {
|
|
segment.SetPKCandidate(bfs)
|
|
bfs.Charge()
|
|
mlog.Info(context.TODO(), "using external real-PK bloom filter candidate",
|
|
mlog.FieldSegmentID(loadInfo.GetSegmentID()))
|
|
}
|
|
if !segment.PkCandidateExist() {
|
|
return merr.WrapErrServiceInternalMsg("milvus-table real-PK segment missing bloom filter stats")
|
|
}
|
|
} else {
|
|
candidate = pkoracle.NewExternalSegmentCandidate(
|
|
loadInfo.GetSegmentID(),
|
|
loadInfo.GetPartitionID(),
|
|
segment.Type(),
|
|
)
|
|
}
|
|
if candidate != nil {
|
|
segment.SetPKCandidate(candidate)
|
|
mlog.Info(context.TODO(), "using external collection PK candidate",
|
|
mlog.FieldSegmentID(loadInfo.GetSegmentID()),
|
|
mlog.Bool("realPK", isMilvusTableRealPK))
|
|
}
|
|
|
|
// Check for truncated segment ID collision with other segments being loaded.
|
|
if !isMilvusTableRealPK {
|
|
collisions := detectVirtualPKCollisions(loadInfo.GetSegmentID(), infos)
|
|
for _, collidingID := range collisions {
|
|
mlog.Warn(context.TODO(), "virtual PK collision detected: two segments share truncated segment ID",
|
|
mlog.Int64("segmentID1", loadInfo.GetSegmentID()),
|
|
mlog.Int64("segmentID2", collidingID),
|
|
mlog.Int64("truncatedID", loadInfo.GetSegmentID()&0xFFFFFFFF))
|
|
}
|
|
}
|
|
} else if paramtable.Get().CommonCfg.BloomFilterEnabled.GetAsBool() {
|
|
bfs, err := loader.loadSingleBloomFilterSet(ctx, loadInfo.GetCollectionID(), loadInfo, segment.Type())
|
|
if err != nil {
|
|
return merr.Wrap(err, "At LoadBloomFilter")
|
|
}
|
|
segment.SetPKCandidate(bfs)
|
|
// Charge bloom filter resource
|
|
bfs.Charge()
|
|
}
|
|
}
|
|
|
|
if segment.Level() != datapb.SegmentLevel_L0 {
|
|
loader.manager.Segment.Put(ctx, segmentType, segment)
|
|
}
|
|
newSegments.GetAndRemove(segmentID)
|
|
loaded.Insert(segmentID, segment)
|
|
loader.notifyLoadFinish(loadInfo)
|
|
|
|
metrics.QueryNodeLoadSegmentLatency.WithLabelValues(paramtable.GetStringNodeID()).Observe(float64(tr.ElapseSpan().Milliseconds()))
|
|
return nil
|
|
}
|
|
|
|
// Start to load,
|
|
// Make sure we can always benefit from concurrency, and not spawn too many idle goroutines
|
|
mlog.Info(context.TODO(), "start to load segments in parallel",
|
|
mlog.Int("segmentNum", len(infos)),
|
|
mlog.Int("concurrencyLevel", requestResourceResult.ConcurrencyLevel))
|
|
|
|
err = funcutil.ProcessFuncParallel(len(infos),
|
|
requestResourceResult.ConcurrencyLevel, loadSegmentFunc, "loadSegmentFunc")
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "failed to load some segments", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
|
|
// Wait for all segments loaded
|
|
segmentIDs := lo.Map(segments, func(info *querypb.SegmentLoadInfo, _ int) int64 { return info.GetSegmentID() })
|
|
if err := loader.waitSegmentLoadDone(ctx, segmentType, segmentIDs, version); err != nil {
|
|
mlog.Warn(context.TODO(), "failed to wait the filtered out segments load done", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
|
|
mlog.Info(context.TODO(), "all segment load done")
|
|
var result []Segment
|
|
loaded.Range(func(_ int64, s Segment) bool {
|
|
result = append(result, s)
|
|
return true
|
|
})
|
|
return result, nil
|
|
}
|
|
|
|
func (loader *segmentLoader) prepare(ctx context.Context, segmentType SegmentType, segments ...*querypb.SegmentLoadInfo) []*querypb.SegmentLoadInfo {
|
|
// filter out loaded & loading segments
|
|
infos := make([]*querypb.SegmentLoadInfo, 0, len(segments))
|
|
for _, segment := range segments {
|
|
// Only active loaded segments should be skipped here. SegmentManager.Exist()
|
|
// also reports detached/on-releasing segments, which are no longer active
|
|
// and must be allowed to load again.
|
|
isLoaded := loader.manager.Segment.GetWithType(segment.GetSegmentID(), segmentType) != nil
|
|
isLoading := loader.loadingSegments.Contain(segment.GetSegmentID())
|
|
if !isLoaded && !isLoading {
|
|
infos = append(infos, segment)
|
|
loader.loadingSegments.Insert(segment.GetSegmentID(), newLoadResult())
|
|
} else {
|
|
mlog.Info(context.TODO(), "skip loaded/loading segment",
|
|
mlog.Int64("segmentID", segment.GetSegmentID()),
|
|
mlog.Bool("isLoaded", isLoaded),
|
|
mlog.Bool("isLoading", isLoading),
|
|
)
|
|
}
|
|
}
|
|
|
|
return infos
|
|
}
|
|
|
|
func configureUseTakeForOutput(loadInfo *querypb.SegmentLoadInfo, schema *schemapb.CollectionSchema) {
|
|
if loadInfo == nil {
|
|
return
|
|
}
|
|
if typeutil.IsExternalCollection(schema) {
|
|
loadInfo.UseTakeForOutput = paramtable.Get().QueryNodeCfg.ExternalCollectionUseTakeForOutput.GetAsBool()
|
|
return
|
|
}
|
|
loadInfo.UseTakeForOutput = paramtable.Get().QueryNodeCfg.InternalCollectionUseTakeForOutput.GetAsBool()
|
|
}
|
|
|
|
func (loader *segmentLoader) unregister(segments ...*querypb.SegmentLoadInfo) {
|
|
for i := range segments {
|
|
result, ok := loader.loadingSegments.GetAndRemove(segments[i].GetSegmentID())
|
|
if ok {
|
|
result.SetResult(failure)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (loader *segmentLoader) notifyLoadFinish(segments ...*querypb.SegmentLoadInfo) {
|
|
for _, loadInfo := range segments {
|
|
result, ok := loader.loadingSegments.Get(loadInfo.GetSegmentID())
|
|
if ok {
|
|
result.SetResult(success)
|
|
}
|
|
}
|
|
}
|
|
|
|
// requestResource requests memory & storage to load segments,
|
|
// returns the memory usage, disk usage and concurrency with the gained memory.
|
|
func (loader *segmentLoader) requestResource(ctx context.Context, infos ...*querypb.SegmentLoadInfo) (requestResourceResult, error) {
|
|
// we need to deal with empty infos case separately,
|
|
// because the following judgement for requested resources are based on current status and static config
|
|
// which may block empty-load operations by accident
|
|
if len(infos) == 0 {
|
|
return requestResourceResult{}, nil
|
|
}
|
|
|
|
segmentIDs := lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) int64 {
|
|
return info.GetSegmentID()
|
|
})
|
|
logger := mlog.With(
|
|
mlog.Int64s("segmentIDs", segmentIDs),
|
|
)
|
|
|
|
loadingUsage, maxSegmentSize, err := loader.estimateSegmentLoadingResourceUsage(ctx, infos...)
|
|
if err != nil {
|
|
logger.Warn(ctx, "no sufficient physical resource to load segments", mlog.Err(err))
|
|
return requestResourceResult{}, err
|
|
}
|
|
|
|
loader.mut.Lock()
|
|
defer loader.mut.Unlock()
|
|
|
|
physicalMemoryUsage := hardware.GetUsedMemoryCount()
|
|
totalMemory := hardware.GetMemoryCount()
|
|
|
|
physicalDiskUsage, err := loader.duf.GetDiskUsage()
|
|
if err != nil {
|
|
return requestResourceResult{}, merr.Wrap(err, "get local used size failed")
|
|
}
|
|
diskCap := paramtable.Get().QueryNodeCfg.DiskCapacityLimit.GetAsUint64()
|
|
|
|
result := requestResourceResult{
|
|
CommittedResource: loader.committedResource,
|
|
}
|
|
|
|
if loader.committedResource.MemorySize+physicalMemoryUsage >= totalMemory {
|
|
return result, merr.WrapErrServiceMemoryLimitExceeded(float32(loader.committedResource.MemorySize+physicalMemoryUsage), float32(totalMemory))
|
|
} else if loader.committedResource.DiskSize+uint64(physicalDiskUsage) >= diskCap {
|
|
return result, merr.WrapErrServiceDiskLimitExceeded(float32(loader.committedResource.DiskSize+uint64(physicalDiskUsage)), float32(diskCap))
|
|
}
|
|
|
|
result.ConcurrencyLevel = funcutil.Min(hardware.GetCPUNum(), len(infos))
|
|
|
|
// TODO: disable logical resource checking for now
|
|
// lmu, ldu, err := loader.checkLogicalSegmentSize(ctx, infos, totalMemory)
|
|
// if err != nil {
|
|
// mlog.Warn(context.TODO(), "no sufficient logical resource to load segments", mlog.Err(err))
|
|
// return result, err
|
|
// }
|
|
|
|
if err := loader.checkLoadingResource(ctx, logger, loadingUsage, maxSegmentSize, totalMemory, physicalMemoryUsage, physicalDiskUsage); err != nil {
|
|
return result, err
|
|
}
|
|
|
|
result.Resource.MemorySize = loadingUsage.MemorySize
|
|
result.Resource.DiskSize = loadingUsage.DiskSize
|
|
// result.LogicalResource.MemorySize = lmu
|
|
// result.LogicalResource.DiskSize = ldu
|
|
|
|
loader.committedResource.Add(result.Resource)
|
|
// loader.committedLogicalResource.Add(result.LogicalResource)
|
|
mlog.Info(context.TODO(), "request resource for loading segments (unit in MiB)",
|
|
mlog.Float64("memory", logutil.ToMB(float64(result.Resource.MemorySize))),
|
|
mlog.Float64("committedMemory", logutil.ToMB(float64(loader.committedResource.MemorySize))),
|
|
mlog.Float64("disk", logutil.ToMB(float64(result.Resource.DiskSize))),
|
|
mlog.Float64("committedDisk", logutil.ToMB(float64(loader.committedResource.DiskSize))),
|
|
)
|
|
|
|
return result, nil
|
|
}
|
|
|
|
// freeRequestResource returns request memory & storage usage request.
|
|
func (loader *segmentLoader) freeRequestResource(requestResourceResult requestResourceResult) {
|
|
loader.mut.Lock()
|
|
defer loader.mut.Unlock()
|
|
|
|
resource := requestResourceResult.Resource
|
|
// logicalResource := requestResourceResult.LogicalResource
|
|
|
|
if paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool() {
|
|
C.ReleaseLoadingResource(C.CResourceUsage{
|
|
memory_bytes: C.int64_t(resource.MemorySize),
|
|
disk_bytes: C.int64_t(resource.DiskSize),
|
|
})
|
|
}
|
|
|
|
loader.committedResource.Sub(resource)
|
|
// loader.committedLogicalResource.Sub(logicalResource)
|
|
loader.committedResourceNotifier.NotifyAll()
|
|
}
|
|
|
|
func (loader *segmentLoader) waitSegmentLoadDone(ctx context.Context, segmentType SegmentType, segmentIDs []int64, version int64) error {
|
|
for _, segmentID := range segmentIDs {
|
|
if loader.manager.Segment.GetWithType(segmentID, segmentType) != nil {
|
|
continue
|
|
}
|
|
|
|
result, ok := loader.loadingSegments.Get(segmentID)
|
|
if !ok {
|
|
mlog.Warn(context.TODO(), "segment was removed from the loading map early", mlog.Int64("segmentID", segmentID))
|
|
return merr.WrapErrServiceInternalMsg("segment was removed from the loading map early")
|
|
}
|
|
|
|
mlog.Info(context.TODO(), "wait segment loaded...", mlog.Int64("segmentID", segmentID))
|
|
|
|
signal := make(chan struct{})
|
|
go func() {
|
|
select {
|
|
case <-signal:
|
|
case <-ctx.Done():
|
|
result.cond.Broadcast()
|
|
}
|
|
}()
|
|
result.cond.L.Lock()
|
|
for result.status.Load() == loading && ctx.Err() == nil {
|
|
result.cond.Wait()
|
|
}
|
|
result.cond.L.Unlock()
|
|
close(signal)
|
|
|
|
if ctx.Err() != nil {
|
|
mlog.Warn(context.TODO(), "failed to wait segment loaded due to context done", mlog.Int64("segmentID", segmentID))
|
|
return ctx.Err()
|
|
}
|
|
|
|
if result.status.Load() == failure {
|
|
mlog.Warn(context.TODO(), "failed to wait segment loaded", mlog.Int64("segmentID", segmentID))
|
|
return merr.WrapErrSegmentLack(segmentID, "failed to wait segment loaded")
|
|
}
|
|
|
|
// try to update segment version after wait segment loaded
|
|
loader.manager.Segment.UpdateBy(IncreaseVersion(version), WithType(segmentType), WithID(segmentID))
|
|
|
|
mlog.Info(context.TODO(), "segment loaded...", mlog.Int64("segmentID", segmentID))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (loader *segmentLoader) GetChunkManager() storage.ChunkManager {
|
|
return loader.cm
|
|
}
|
|
|
|
// load single bloom filter
|
|
func (loader *segmentLoader) loadSingleBloomFilterSet(ctx context.Context, collectionID int64, loadInfo *querypb.SegmentLoadInfo, segtype SegmentType) (*pkoracle.BloomFilterSet, error) {
|
|
partitionID := loadInfo.PartitionID
|
|
segmentID := loadInfo.SegmentID
|
|
bfs := pkoracle.NewBloomFilterSet(segmentID, partitionID, segtype)
|
|
|
|
collection := loader.manager.Collection.Get(collectionID)
|
|
if collection == nil {
|
|
err := merr.WrapErrCollectionNotFound(collectionID)
|
|
mlog.Warn(context.TODO(), "failed to get collection while loading segment", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
|
|
mlog.Info(context.TODO(), "start loading remote...", mlog.Int("segmentNum", 1))
|
|
|
|
schema := collection.Schema()
|
|
isExternalCollection := typeutil.IsExternalCollection(schema)
|
|
isMilvusTableRealPK := typeutil.NewStorageColumnResolver(schema).IsMilvusTable() &&
|
|
HasExternalPrimaryKey(schema)
|
|
if !paramtable.Get().CommonCfg.BloomFilterEnabled.GetAsBool() && !isMilvusTableRealPK {
|
|
mlog.Info(context.TODO(), "skip loading bloom filter for remote segment because bloom filter is disabled")
|
|
return bfs, nil
|
|
}
|
|
if isExternalCollection && !isMilvusTableRealPK {
|
|
mlog.Debug(context.TODO(), "virtual-PK external collection: returning empty bloom filter set")
|
|
return bfs, nil
|
|
}
|
|
|
|
pkField := GetPkField(schema)
|
|
mlog.Info(context.TODO(), "loading bloom filter for remote...")
|
|
pkStatsBinlogs, err := packed.NewStatsResolverFromLoadInfo(loadInfo).BloomFilterPaths(pkField.GetFieldID())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
err = loader.loadBloomFilter(ctx, segmentID, bfs, pkStatsBinlogs, loader.bloomFilterDownloader(collection, isMilvusTableRealPK))
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "load remote segment bloom filter failed",
|
|
mlog.Int64("partitionID", partitionID),
|
|
mlog.Int64("segmentID", segmentID),
|
|
mlog.Err(err),
|
|
)
|
|
return nil, err
|
|
}
|
|
if isMilvusTableRealPK && !bfs.PkCandidateExist() {
|
|
return nil, merr.WrapErrServiceInternalMsg("milvus-table real-PK segment missing bloom filter stats")
|
|
}
|
|
|
|
return bfs, nil
|
|
}
|
|
|
|
func (loader *segmentLoader) LoadBloomFilterSet(ctx context.Context, collectionID int64, infos ...*querypb.SegmentLoadInfo) ([]*pkoracle.BloomFilterSet, error) {
|
|
segmentNum := len(infos)
|
|
if segmentNum == 0 {
|
|
mlog.Info(context.TODO(), "no segment to load")
|
|
return nil, nil
|
|
}
|
|
|
|
// Phase 1: always create metadata-only stubs (segmentID / partitionID / type).
|
|
// This gives callers valid candidates even when BF data is not loaded,
|
|
// so partition filtering and type-based delete-scope logic never need nil guards.
|
|
bfSets := make([]*pkoracle.BloomFilterSet, segmentNum)
|
|
for i, info := range infos {
|
|
bfSets[i] = pkoracle.NewBloomFilterSet(info.GetSegmentID(), info.GetPartitionID(), commonpb.SegmentState_Sealed)
|
|
}
|
|
|
|
collection := loader.manager.Collection.Get(collectionID)
|
|
if collection == nil {
|
|
err := merr.WrapErrCollectionNotFound(collectionID)
|
|
mlog.Warn(context.TODO(), "failed to get collection while loading segment", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
|
|
schema := collection.Schema()
|
|
isExternalCollection := typeutil.IsExternalCollection(schema)
|
|
isMilvusTableRealPK := typeutil.NewStorageColumnResolver(schema).IsMilvusTable() &&
|
|
HasExternalPrimaryKey(schema)
|
|
|
|
// Phase 2: load BF stats into the stubs. Milvus-table real-PK correctness
|
|
// depends on source bloom filters, so that path ignores the global BF
|
|
// disable switch; other collections keep the historical metadata-only
|
|
// behavior when BloomFilterEnabled=false.
|
|
if !paramtable.Get().CommonCfg.BloomFilterEnabled.GetAsBool() && !isMilvusTableRealPK {
|
|
mlog.Info(context.TODO(), "bloom filter disabled: returning metadata-only stubs")
|
|
return bfSets, nil
|
|
}
|
|
|
|
// Virtual-PK external collections use ExternalSegmentCandidate and have no
|
|
// reusable source-side PK stats.
|
|
if isExternalCollection && !isMilvusTableRealPK {
|
|
return bfSets, nil
|
|
}
|
|
|
|
pkField := GetPkField(schema)
|
|
pkFieldID := pkField.GetFieldID()
|
|
|
|
// Calculate total memory size needed for bloom filters (PK stats)
|
|
var totalMemorySize int64
|
|
for _, info := range infos {
|
|
memSize, _ := packed.NewStatsResolverFromLoadInfo(info).BloomFilterMemorySize(pkFieldID)
|
|
totalMemorySize += memSize
|
|
}
|
|
|
|
// Reserve memory resource if tiered eviction is enabled
|
|
if paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool() && totalMemorySize > 0 {
|
|
if ok := C.TryReserveLoadingResourceWithTimeout(C.CResourceUsage{
|
|
// double loading memory size for bloom filters to avoid OOM during loading
|
|
memory_bytes: C.int64_t(totalMemorySize * 2),
|
|
disk_bytes: C.int64_t(0),
|
|
}, 1000); !ok {
|
|
return nil, merr.WrapErrSegmentRequestResourceFailed("memory",
|
|
fmt.Sprintf("failed to reserve loading resource for bloom filters, totalMemorySize = %v MB",
|
|
logutil.ToMB(float64(totalMemorySize))))
|
|
}
|
|
mlog.Info(context.TODO(), "reserved loading resource for bloom filters", mlog.Float64("totalMemorySizeMB", logutil.ToMB(float64(totalMemorySize))))
|
|
}
|
|
|
|
defer func() {
|
|
if paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool() && totalMemorySize > 0 {
|
|
C.ReleaseLoadingResource(C.CResourceUsage{
|
|
memory_bytes: C.int64_t(totalMemorySize * 2),
|
|
disk_bytes: C.int64_t(0),
|
|
})
|
|
mlog.Info(context.TODO(), "released loading resource for bloom filters", mlog.Float64("totalMemorySizeMB", logutil.ToMB(float64(totalMemorySize))))
|
|
}
|
|
}()
|
|
|
|
mlog.Info(context.TODO(), "start loading remote...", mlog.Int("segmentNum", segmentNum))
|
|
|
|
loadRemoteFunc := func(idx int) error {
|
|
loadInfo := infos[idx]
|
|
bfs := bfSets[idx]
|
|
|
|
mlog.Info(context.TODO(), "loading bloom filter for remote...")
|
|
pkStatsBinlogs, err := packed.NewStatsResolverFromLoadInfo(loadInfo).BloomFilterPaths(pkFieldID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
err = loader.loadBloomFilter(ctx, bfs.ID(), bfs, pkStatsBinlogs, loader.bloomFilterDownloader(collection, isMilvusTableRealPK))
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "load remote segment bloom filter failed",
|
|
mlog.Int64("partitionID", bfs.Partition()),
|
|
mlog.Int64("segmentID", bfs.ID()),
|
|
mlog.Err(err),
|
|
)
|
|
return err
|
|
}
|
|
if isMilvusTableRealPK && !bfs.PkCandidateExist() {
|
|
return merr.WrapErrServiceInternalMsg("milvus-table real-PK segment missing bloom filter stats")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
err := funcutil.ProcessFuncParallel(segmentNum, segmentNum, loadRemoteFunc, "loadRemoteFunc")
|
|
if err != nil {
|
|
// no partial success here
|
|
mlog.Warn(context.TODO(), "failed to load remote segment", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
|
|
// Charge loaded resource for bloom filters
|
|
for _, bfs := range bfSets {
|
|
bfs.Charge()
|
|
}
|
|
|
|
return bfSets, nil
|
|
}
|
|
|
|
func separateIndexAndBinlog(loadInfo *querypb.SegmentLoadInfo) (map[int64]*IndexedFieldInfo, []*datapb.FieldBinlog) {
|
|
fieldID2IndexInfo := make(map[int64][]*querypb.FieldIndexInfo)
|
|
for _, indexInfo := range loadInfo.IndexInfos {
|
|
if len(indexInfo.GetIndexFilePaths()) > 0 {
|
|
fieldID := indexInfo.FieldID
|
|
fieldID2IndexInfo[fieldID] = append(fieldID2IndexInfo[fieldID], indexInfo)
|
|
}
|
|
}
|
|
|
|
preferFieldData := paramtable.Get().QueryNodeCfg.PreferFieldDataWhenIndexHasRawData.GetAsBool()
|
|
|
|
indexedFieldInfos := make(map[int64]*IndexedFieldInfo)
|
|
fieldBinlogs := make([]*datapb.FieldBinlog, 0, len(loadInfo.BinlogPaths))
|
|
|
|
for _, fieldBinlog := range loadInfo.BinlogPaths {
|
|
fieldID := fieldBinlog.FieldID
|
|
// check num rows of data meta and index meta are consistent
|
|
if indexInfo, ok := fieldID2IndexInfo[fieldID]; ok {
|
|
for _, index := range indexInfo {
|
|
fieldInfo := &IndexedFieldInfo{
|
|
FieldBinlog: fieldBinlog,
|
|
IndexInfo: index,
|
|
}
|
|
indexedFieldInfos[index.IndexID] = fieldInfo
|
|
}
|
|
if preferFieldData {
|
|
fieldBinlogs = append(fieldBinlogs, fieldBinlog)
|
|
}
|
|
} else {
|
|
fieldBinlogs = append(fieldBinlogs, fieldBinlog)
|
|
}
|
|
}
|
|
|
|
return indexedFieldInfos, fieldBinlogs
|
|
}
|
|
|
|
// detectVirtualPKCollisions checks if any segments in infos share the same
|
|
// truncated (lower 32 bits) segment ID as segmentID. A collision means two
|
|
// segments produce overlapping virtual PK spaces.
|
|
func detectVirtualPKCollisions(segmentID int64, infos []*querypb.SegmentLoadInfo) []int64 {
|
|
truncatedID := segmentID & 0xFFFFFFFF
|
|
var collisions []int64
|
|
for _, info := range infos {
|
|
if info.GetSegmentID() != segmentID &&
|
|
(info.GetSegmentID()&0xFFFFFFFF) == truncatedID {
|
|
collisions = append(collisions, info.GetSegmentID())
|
|
}
|
|
}
|
|
return collisions
|
|
}
|
|
|
|
func separateLoadInfoV2(loadInfo *querypb.SegmentLoadInfo, schema *schemapb.CollectionSchema) (
|
|
map[int64]*IndexedFieldInfo, // indexed info
|
|
[]*datapb.FieldBinlog, // fields info
|
|
map[int64]*datapb.TextIndexStats, // text indexed info
|
|
map[int64]struct{}, // unindexed text fields
|
|
map[int64]*datapb.JsonKeyStats, // json key stats info
|
|
map[int64]string, // text index base paths
|
|
map[int64]string, // json key stats base paths
|
|
) {
|
|
storageVersion := loadInfo.GetStorageVersion()
|
|
|
|
// Build a map of external field IDs for quick lookup
|
|
// External fields are skipped during loading (lazy loaded on demand)
|
|
externalFieldIDs := make(map[int64]bool)
|
|
isExternalColl := typeutil.IsExternalCollection(schema)
|
|
if isExternalColl {
|
|
for _, field := range schema.GetFields() {
|
|
if IsExternalField(field) {
|
|
externalFieldIDs[field.GetFieldID()] = true
|
|
}
|
|
}
|
|
}
|
|
|
|
fieldID2IndexInfo := make(map[int64][]*querypb.FieldIndexInfo)
|
|
for _, indexInfo := range loadInfo.IndexInfos {
|
|
if len(indexInfo.GetIndexFilePaths()) > 0 {
|
|
fieldID := indexInfo.FieldID
|
|
fieldID2IndexInfo[fieldID] = append(fieldID2IndexInfo[fieldID], indexInfo)
|
|
}
|
|
}
|
|
|
|
preferFieldData := paramtable.Get().QueryNodeCfg.PreferFieldDataWhenIndexHasRawData.GetAsBool()
|
|
|
|
indexedFieldInfos := make(map[int64]*IndexedFieldInfo)
|
|
fieldBinlogs := make([]*datapb.FieldBinlog, 0, len(loadInfo.BinlogPaths))
|
|
|
|
if storageVersion == storage.StorageV2 || storageVersion == storage.StorageV3 {
|
|
for _, fieldBinlog := range loadInfo.BinlogPaths {
|
|
fieldID := fieldBinlog.FieldID
|
|
|
|
// Skip external fields - they are lazy loaded on demand
|
|
if externalFieldIDs[fieldID] {
|
|
continue
|
|
}
|
|
|
|
if fieldID == storagecommon.DefaultShortColumnGroupID {
|
|
allFields := typeutil.GetAllFieldSchemas(schema)
|
|
// for short column group, we need to load all fields in the group
|
|
for _, field := range allFields {
|
|
// Skip external fields in short column group
|
|
if externalFieldIDs[field.GetFieldID()] {
|
|
continue
|
|
}
|
|
if infos, ok := fieldID2IndexInfo[field.GetFieldID()]; ok {
|
|
for _, indexInfo := range infos {
|
|
fieldInfo := &IndexedFieldInfo{
|
|
FieldBinlog: fieldBinlog,
|
|
IndexInfo: indexInfo,
|
|
}
|
|
indexedFieldInfos[indexInfo.IndexID] = fieldInfo
|
|
}
|
|
}
|
|
}
|
|
fieldBinlogs = append(fieldBinlogs, fieldBinlog)
|
|
} else {
|
|
// for single file field, such as vector field, text field
|
|
if infos, ok := fieldID2IndexInfo[fieldID]; ok {
|
|
for _, indexInfo := range infos {
|
|
fieldInfo := &IndexedFieldInfo{
|
|
FieldBinlog: fieldBinlog,
|
|
IndexInfo: indexInfo,
|
|
}
|
|
indexedFieldInfos[indexInfo.IndexID] = fieldInfo
|
|
}
|
|
if preferFieldData {
|
|
fieldBinlogs = append(fieldBinlogs, fieldBinlog)
|
|
}
|
|
} else {
|
|
fieldBinlogs = append(fieldBinlogs, fieldBinlog)
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
for _, fieldBinlog := range loadInfo.BinlogPaths {
|
|
fieldID := fieldBinlog.FieldID
|
|
|
|
// Skip external fields - they are lazy loaded on demand
|
|
if externalFieldIDs[fieldID] {
|
|
continue
|
|
}
|
|
|
|
if infos, ok := fieldID2IndexInfo[fieldID]; ok {
|
|
for _, indexInfo := range infos {
|
|
fieldInfo := &IndexedFieldInfo{
|
|
FieldBinlog: fieldBinlog,
|
|
IndexInfo: indexInfo,
|
|
}
|
|
indexedFieldInfos[indexInfo.IndexID] = fieldInfo
|
|
}
|
|
if preferFieldData {
|
|
fieldBinlogs = append(fieldBinlogs, fieldBinlog)
|
|
}
|
|
} else {
|
|
fieldBinlogs = append(fieldBinlogs, fieldBinlog)
|
|
}
|
|
}
|
|
}
|
|
|
|
// For external table segments (ManifestPath set, BinlogPaths empty), extract
|
|
// indexes directly from fieldID2IndexInfo without a corresponding FieldBinlog,
|
|
// because the segment data lives in the external store rather than Milvus binlogs.
|
|
if loadInfo.GetManifestPath() != "" {
|
|
for _, infos := range fieldID2IndexInfo {
|
|
for _, indexInfo := range infos {
|
|
if _, exists := indexedFieldInfos[indexInfo.IndexID]; !exists {
|
|
indexedFieldInfos[indexInfo.IndexID] = &IndexedFieldInfo{
|
|
FieldBinlog: &datapb.FieldBinlog{},
|
|
IndexInfo: indexInfo,
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
statsResult := packed.NewStatsResolverFromLoadInfo(loadInfo).TextAndJSONIndexStatsWithBasePaths()
|
|
textIndexedInfo := statsResult.TextIndexStats
|
|
jsonKeyIndexInfo := statsResult.JSONKeyStats
|
|
textBasePaths := statsResult.TextBasePaths
|
|
jsonBasePaths := statsResult.JSONBasePaths
|
|
if statsResult.Err() != nil {
|
|
mlog.Warn(context.TODO(), "failed to load text/json stats from manifest",
|
|
mlog.String("manifestPath", loadInfo.GetManifestPath()), mlog.Err(statsResult.Err()))
|
|
textIndexedInfo = make(map[int64]*datapb.TextIndexStats)
|
|
jsonKeyIndexInfo = make(map[int64]*datapb.JsonKeyStats)
|
|
textBasePaths = make(map[int64]string)
|
|
jsonBasePaths = make(map[int64]string)
|
|
}
|
|
|
|
if textBasePaths == nil {
|
|
textBasePaths = make(map[int64]string)
|
|
}
|
|
if jsonBasePaths == nil {
|
|
jsonBasePaths = make(map[int64]string)
|
|
}
|
|
|
|
// For V2 (non-manifest) segments, compute basePaths from metadata.
|
|
// The resolver returns empty basePaths for V2; we compute them here.
|
|
rootPath := paramtable.Get().MinioCfg.RootPath.GetValue()
|
|
for fieldID, stats := range textIndexedInfo {
|
|
if _, ok := textBasePaths[fieldID]; !ok {
|
|
textBasePaths[fieldID] = metautil.BuildTextIndexPrefix(rootPath,
|
|
stats.GetBuildID(), stats.GetVersion(),
|
|
loadInfo.GetCollectionID(), loadInfo.GetPartitionID(), loadInfo.GetSegmentID(), fieldID)
|
|
}
|
|
}
|
|
for fieldID, stats := range jsonKeyIndexInfo {
|
|
if _, ok := jsonBasePaths[fieldID]; !ok {
|
|
jsonBasePaths[fieldID] = metautil.BuildJSONKeyStatsPrefix(rootPath, stats.GetJsonKeyStatsDataFormat(),
|
|
stats.GetBuildID(), stats.GetVersion(),
|
|
loadInfo.GetCollectionID(), loadInfo.GetPartitionID(), loadInfo.GetSegmentID(), fieldID)
|
|
}
|
|
}
|
|
|
|
unindexedTextFields := make(map[int64]struct{})
|
|
// todo(SpadeA): consider struct fields when index is ready
|
|
for _, field := range schema.GetFields() {
|
|
h := typeutil.CreateFieldSchemaHelper(field)
|
|
_, textIndexExist := textIndexedInfo[field.GetFieldID()]
|
|
if h.EnableMatch() && !textIndexExist {
|
|
unindexedTextFields[field.GetFieldID()] = struct{}{}
|
|
}
|
|
}
|
|
|
|
return indexedFieldInfos, fieldBinlogs, textIndexedInfo, unindexedTextFields, jsonKeyIndexInfo, textBasePaths, jsonBasePaths
|
|
}
|
|
|
|
func (loader *segmentLoader) loadSealedSegment(ctx context.Context, loadInfo *querypb.SegmentLoadInfo, segment *LocalSegment) (err error) {
|
|
// TODO: we should create a transaction-like api to load segment for segment interface,
|
|
// but not do many things in segment loader.
|
|
stateLockGuard, err := segment.StartLoadData()
|
|
// segment can not do load now.
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if stateLockGuard == nil {
|
|
return nil
|
|
}
|
|
defer func() {
|
|
if err != nil {
|
|
// Release partial loaded segment data if load failed.
|
|
segment.ReleaseSegmentData()
|
|
}
|
|
stateLockGuard.Done(err)
|
|
}()
|
|
|
|
collection := segment.GetCollection()
|
|
indexedFieldInfos, _, textIndexes, unindexedTextFields, jsonKeyStats, _, _ := separateLoadInfoV2(loadInfo, collection.Schema())
|
|
|
|
tr := timerecord.NewTimeRecorder("segmentLoader.loadSealedSegment")
|
|
mlog.Info(context.TODO(), "Start loading fields...",
|
|
mlog.Int("indexedFields count", len(indexedFieldInfos)),
|
|
mlog.Int64s("indexed text fields", lo.Keys(textIndexes)),
|
|
mlog.Int64s("unindexed text fields", lo.Keys(unindexedTextFields)),
|
|
mlog.Int64s("indexed json key fields", lo.Keys(jsonKeyStats)),
|
|
)
|
|
_, err = GetLoadPool().Submit(func() (any, error) {
|
|
if err = segment.Load(ctx); err != nil {
|
|
return struct{}{}, merr.Wrap(err, "At Load")
|
|
}
|
|
|
|
return struct{}{}, nil
|
|
}).Await()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, indexInfo := range loadInfo.IndexInfos {
|
|
segment.fieldIndexes.Insert(indexInfo.GetIndexID(), &IndexedFieldInfo{
|
|
FieldBinlog: &datapb.FieldBinlog{
|
|
FieldID: indexInfo.GetFieldID(),
|
|
},
|
|
IndexInfo: indexInfo,
|
|
IsLoaded: true,
|
|
})
|
|
}
|
|
|
|
// 4. rectify entries number for binlog in very rare cases
|
|
// https://github.com/milvus-io/milvus/23654
|
|
// legacy entry num = 0
|
|
if err := loader.patchEntryNumber(ctx, segment, loadInfo); err != nil {
|
|
return err
|
|
}
|
|
patchEntryNumberSpan := tr.RecordSpan()
|
|
mlog.Info(context.TODO(), "Finish loading segment",
|
|
mlog.Duration("patchEntryNumberSpan", patchEntryNumberSpan),
|
|
)
|
|
return nil
|
|
}
|
|
|
|
func (loader *segmentLoader) LoadSegment(ctx context.Context,
|
|
seg Segment,
|
|
loadInfo *querypb.SegmentLoadInfo,
|
|
) (err error) {
|
|
segment, ok := seg.(*LocalSegment)
|
|
if !ok {
|
|
return merr.WrapErrParameterInvalid("LocalSegment", fmt.Sprintf("%T", seg))
|
|
}
|
|
mlog.Info(context.TODO(), "start loading segment files",
|
|
mlog.Int64("rowNum", loadInfo.GetNumOfRows()),
|
|
mlog.String("segmentType", segment.Type().String()),
|
|
mlog.Int32("priority", int32(loadInfo.GetPriority())))
|
|
|
|
collection := loader.manager.Collection.Get(segment.Collection())
|
|
if collection == nil {
|
|
err := merr.WrapErrCollectionNotFound(segment.Collection())
|
|
mlog.Warn(context.TODO(), "failed to get collection while loading segment", mlog.Err(err))
|
|
return err
|
|
}
|
|
pkField := GetPkField(collection.Schema())
|
|
|
|
if segment.Type() == SegmentTypeSealed {
|
|
if err := loader.loadSealedSegment(ctx, loadInfo, segment); err != nil {
|
|
return err
|
|
}
|
|
} else {
|
|
if err := segment.Load(ctx); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
binlogSize := calculateSegmentMemorySize(segment.LoadInfo())
|
|
segment.manager.AddLoadedBinlogSize(binlogSize)
|
|
segment.binlogSize.Store(binlogSize)
|
|
|
|
// load statslog if it's growing segment
|
|
if segment.segmentType == SegmentTypeGrowing {
|
|
if bf, ok := segment.pkCandidate.(*pkoracle.BloomFilterSet); ok {
|
|
mlog.Info(context.TODO(), "loading statslog...")
|
|
resolver := packed.NewStatsResolverFromLoadInfo(loadInfo)
|
|
bfPaths, err := resolver.BloomFilterPaths(pkField.GetFieldID())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := loader.loadBloomFilter(ctx, segment.ID(), bf, bfPaths, loader.cm.MultiRead); err != nil {
|
|
return err
|
|
}
|
|
|
|
bm25Paths, err := resolver.BM25StatsPaths()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
bm25Stats := make(map[int64]*storage.BM25Stats)
|
|
if err := loader.loadBm25Stats(ctx, segment.ID(), bm25Stats, bm25Paths); err != nil {
|
|
return err
|
|
}
|
|
segment.UpdateBM25Stats(bm25Stats)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func loadSealedSegmentFields(ctx context.Context, collection *Collection, segment *LocalSegment, fields []*datapb.FieldBinlog, rowCount int64) error {
|
|
runningGroup, _ := errgroup.WithContext(ctx)
|
|
for _, field := range fields {
|
|
fieldBinLog := field
|
|
fieldID := field.FieldID
|
|
runningGroup.Go(func() error {
|
|
return segment.LoadFieldData(ctx, fieldID, rowCount, fieldBinLog)
|
|
})
|
|
}
|
|
err := runningGroup.Wait()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
mlog.Info(ctx, "load field binlogs done for sealed segment",
|
|
mlog.Int64("collection", segment.Collection()),
|
|
mlog.Int64("segment", segment.ID()),
|
|
mlog.Int("len(field)", len(fields)),
|
|
mlog.String("segmentType", segment.Type().String()))
|
|
|
|
return nil
|
|
}
|
|
|
|
func (loader *segmentLoader) loadFieldsIndex(ctx context.Context,
|
|
schemaHelper *typeutil.SchemaHelper,
|
|
segment *LocalSegment,
|
|
numRows int64,
|
|
indexedFieldInfos map[int64]*IndexedFieldInfo,
|
|
) error {
|
|
for _, fieldInfo := range indexedFieldInfos {
|
|
fieldID := fieldInfo.IndexInfo.FieldID
|
|
indexInfo := fieldInfo.IndexInfo
|
|
tr := timerecord.NewTimeRecorder("loadFieldIndex")
|
|
err := loader.loadFieldIndex(ctx, segment, indexInfo)
|
|
loadFieldIndexSpan := tr.RecordSpan()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
mlog.Info(context.TODO(), "load field binlogs done for sealed segment with index",
|
|
mlog.Int64("fieldID", fieldID),
|
|
mlog.Any("binlog", fieldInfo.FieldBinlog.Binlogs),
|
|
mlog.Int32("current_index_version", fieldInfo.IndexInfo.GetCurrentIndexVersion()),
|
|
mlog.Duration("load_duration", loadFieldIndexSpan),
|
|
)
|
|
|
|
// set average row data size of variable field
|
|
field, err := schemaHelper.GetFieldFromID(fieldID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if typeutil.IsVariableDataType(field.GetDataType()) {
|
|
err = segment.UpdateFieldRawDataSize(ctx, numRows, fieldInfo.FieldBinlog)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (loader *segmentLoader) loadBm25Stats(ctx context.Context, segmentID int64, stats map[int64]*storage.BM25Stats, binlogPaths map[int64][]string) error {
|
|
if len(binlogPaths) == 0 {
|
|
mlog.Info(context.TODO(), "there are no bm25 stats logs saved with segment")
|
|
return nil
|
|
}
|
|
|
|
pathList := []string{}
|
|
fieldList := []int64{}
|
|
fieldOffset := []int{}
|
|
for fieldId, logpaths := range binlogPaths {
|
|
pathList = append(pathList, logpaths...)
|
|
fieldList = append(fieldList, fieldId)
|
|
fieldOffset = append(fieldOffset, len(logpaths))
|
|
}
|
|
|
|
startTs := time.Now()
|
|
values, err := loader.cm.MultiRead(ctx, pathList)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
cnt := 0
|
|
for i, fieldID := range fieldList {
|
|
newStats, ok := stats[fieldID]
|
|
if !ok {
|
|
newStats = storage.NewBM25Stats()
|
|
stats[fieldID] = newStats
|
|
}
|
|
|
|
for j := 0; j < fieldOffset[i]; j++ {
|
|
err := newStats.Deserialize(values[cnt+j])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
cnt += fieldOffset[i]
|
|
mlog.Info(context.TODO(), "Successfully load bm25 stats", mlog.Duration("time", time.Since(startTs)), mlog.Int64("numRow", newStats.NumRow()), mlog.Int64("fieldID", fieldID))
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (loader *segmentLoader) loadFieldIndex(ctx context.Context, segment *LocalSegment, indexInfo *querypb.FieldIndexInfo) error {
|
|
filteredPaths := make([]string, 0, len(indexInfo.IndexFilePaths))
|
|
|
|
for _, indexPath := range indexInfo.IndexFilePaths {
|
|
if path.Base(indexPath) != storage.IndexParamsKey {
|
|
filteredPaths = append(filteredPaths, indexPath)
|
|
}
|
|
}
|
|
|
|
indexInfo.IndexFilePaths = filteredPaths
|
|
fieldType, err := loader.getFieldType(segment.Collection(), indexInfo.FieldID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
collection := loader.manager.Collection.Get(segment.Collection())
|
|
if collection == nil {
|
|
return merr.WrapErrCollectionNotLoaded(segment.Collection(), "failed to load field index")
|
|
}
|
|
|
|
return segment.LoadIndex(ctx, indexInfo, fieldType)
|
|
}
|
|
|
|
func (loader *segmentLoader) loadBloomFilter(
|
|
ctx context.Context,
|
|
segmentID int64,
|
|
bfs *pkoracle.BloomFilterSet,
|
|
binlogPaths []string,
|
|
downloader func(context.Context, []string) ([][]byte, error),
|
|
) error {
|
|
return loader.loadBloomFilterWithDownloader(ctx, segmentID, bfs, binlogPaths, downloader)
|
|
}
|
|
|
|
// bloomFilterDownloader returns the byte downloader used for PK bloom-filter
|
|
// stats. Milvus-table real-PK stats live in the external source filesystem;
|
|
// ordinary internal stats stay on the local chunk manager.
|
|
func (loader *segmentLoader) bloomFilterDownloader(
|
|
collection *Collection,
|
|
useExternalSpec bool,
|
|
) func(context.Context, []string) ([][]byte, error) {
|
|
if !useExternalSpec {
|
|
return loader.cm.MultiRead
|
|
}
|
|
schema := collection.Schema()
|
|
extfs := packed.ExternalSpecContext{
|
|
CollectionID: collection.ID(),
|
|
Source: schema.GetExternalSource(),
|
|
Spec: schema.GetExternalSpec(),
|
|
}
|
|
return func(ctx context.Context, paths []string) ([][]byte, error) {
|
|
return readExternalFiles(ctx, createStorageConfig(), extfs, paths)
|
|
}
|
|
}
|
|
|
|
// loadBloomFilterWithDownloader merges one or more serialized PK bloom-filter
|
|
// stats into the segment BloomFilterSet.
|
|
func (loader *segmentLoader) loadBloomFilterWithDownloader(
|
|
ctx context.Context,
|
|
segmentID int64,
|
|
bfs *pkoracle.BloomFilterSet,
|
|
binlogPaths []string,
|
|
downloader func(context.Context, []string) ([][]byte, error),
|
|
) error {
|
|
if len(binlogPaths) == 0 {
|
|
mlog.Info(context.TODO(), "there are no stats logs saved with segment")
|
|
return nil
|
|
}
|
|
|
|
startTs := time.Now()
|
|
values, err := downloader(ctx, binlogPaths)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
blobs := make([]*storage.Blob, len(values))
|
|
for i := range values {
|
|
blobs[i] = &storage.Blob{Value: values[i]}
|
|
}
|
|
|
|
stats, err := storage.DeserializeBloomFilterStats(binlogPaths, blobs)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "failed to deserialize bloom filter stats", mlog.Err(err))
|
|
return err
|
|
}
|
|
|
|
var size uint
|
|
for _, stat := range stats {
|
|
pkStat := &storage.PkStatistics{
|
|
PkFilter: stat.BF,
|
|
MinPK: stat.MinPk,
|
|
MaxPK: stat.MaxPk,
|
|
}
|
|
size += stat.BF.Cap()
|
|
bfs.AddHistoricalStats(pkStat)
|
|
}
|
|
mlog.Info(context.TODO(), "Successfully load pk stats", mlog.Duration("time", time.Since(startTs)), mlog.Uint("size", size))
|
|
return nil
|
|
}
|
|
|
|
// loadDeltalogs performs the internal actions of `LoadDeltaLogs`
|
|
// this function does not perform resource check and is meant be used among other load APIs.
|
|
func (loader *segmentLoader) loadDeltalogs(ctx context.Context, segment Segment, loadInfo *querypb.SegmentLoadInfo) error {
|
|
deltaLogs := loadInfo.GetDeltalogs()
|
|
ctx, sp := otel.Tracer(typeutil.QueryNodeRole).Start(ctx, fmt.Sprintf("LoadDeltalogs-%d", segment.ID()))
|
|
defer sp.End()
|
|
mlog.Info(context.TODO(), "loading delta...")
|
|
|
|
var rowNums int64
|
|
valid := func(binlog *datapb.Binlog, _ int) bool {
|
|
// the segment has applied the delta logs, skip it
|
|
if binlog.GetTimestampTo() > 0 && // this field may be missed in legacy versions
|
|
binlog.GetTimestampTo() < segment.LastDeltaTimestamp() {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
for _, deltaLog := range deltaLogs {
|
|
rowNums += lo.SumBy(lo.Filter(deltaLog.GetBinlogs(), valid), func(binlog *datapb.Binlog) int64 {
|
|
return binlog.GetEntriesNum()
|
|
})
|
|
}
|
|
|
|
collection := loader.manager.Collection.Get(segment.Collection())
|
|
|
|
helper, _ := typeutil.CreateSchemaHelper(collection.Schema())
|
|
pkField, _ := helper.GetPrimaryKeyField()
|
|
deltaData, err := storage.NewDeltaDataWithPkType(rowNums, pkField.DataType)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
readDeltaRecords := func(reader storage.RecordReader) error {
|
|
defer reader.Close()
|
|
for {
|
|
dl, err := reader.Next()
|
|
if err != nil {
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
return err
|
|
}
|
|
|
|
for i := 0; i < dl.Len(); i++ {
|
|
var pk storage.PrimaryKey
|
|
switch pkField.DataType {
|
|
case schemapb.DataType_Int64:
|
|
pk = storage.NewInt64PrimaryKey(dl.Column(0).(*array.Int64).Value(i))
|
|
case schemapb.DataType_VarChar:
|
|
pk = storage.NewVarCharPrimaryKey(dl.Column(0).(*array.String).Value(i))
|
|
}
|
|
ts := typeutil.Timestamp(dl.Column(1).(*array.Int64).Value(i))
|
|
err = deltaData.Append(pk, ts)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
schema := collection.Schema()
|
|
isExternalCollection := typeutil.IsExternalCollection(schema)
|
|
resolver := typeutil.NewStorageColumnResolver(schema)
|
|
if isExternalCollection && !resolver.IsMilvusTable() {
|
|
mlog.Info(context.TODO(), "skip loading delta logs for non-milvus-table external collection")
|
|
return nil
|
|
}
|
|
isMilvusTableRealPK := resolver.IsMilvusTable() && HasExternalPrimaryKey(schema)
|
|
useExplicitDeltalogs := isMilvusTableRealPK && len(deltaLogs) > 0
|
|
readPaths := func(paths []string, opts ...storage.RwOption) error {
|
|
if len(paths) == 0 {
|
|
return nil
|
|
}
|
|
reader, err := storage.NewDeltalogReader(pkField.DataType, paths, opts...)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return readDeltaRecords(reader)
|
|
}
|
|
|
|
if manifestPath := loadInfo.GetManifestPath(); manifestPath != "" && !useExplicitDeltalogs {
|
|
if isMilvusTableRealPK {
|
|
// Real-PK milvus-table manifests keep source deltalogs. Target-owned
|
|
// deltalogs are only valid for virtual-PK translation.
|
|
extfs := packed.ExternalSpecContext{
|
|
CollectionID: collection.ID(),
|
|
Source: schema.GetExternalSource(),
|
|
Spec: schema.GetExternalSpec(),
|
|
}
|
|
sourceDeltalogs, err := packed.GetDeltaLogsFromManifestWithExtfs(
|
|
manifestPath,
|
|
createStorageConfig(),
|
|
extfs,
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := validateMilvusTableRealPKDeltalogPaths(manifestPath, milvusTableDeltalogPaths(sourceDeltalogs)); err != nil {
|
|
return err
|
|
}
|
|
if len(sourceDeltalogs) > 0 {
|
|
storageV3Paths := make([]string, 0)
|
|
legacyPaths := make([]string, 0)
|
|
for _, deltalog := range sourceDeltalogs {
|
|
for _, binlog := range lo.Filter(deltalog.GetBinlogs(), valid) {
|
|
if packed.IsMilvusTableStorageV3DeltalogPath(binlog.GetLogPath()) {
|
|
storageV3Paths = append(storageV3Paths, binlog.GetLogPath())
|
|
} else {
|
|
legacyPaths = append(legacyPaths, binlog.GetLogPath())
|
|
}
|
|
}
|
|
}
|
|
if len(storageV3Paths) > 0 {
|
|
reader, err := storage.NewDeltalogReader(
|
|
pkField.DataType,
|
|
storageV3Paths,
|
|
storage.WithVersion(storage.StorageV3),
|
|
storage.WithStorageConfig(createStorageConfig()),
|
|
storage.WithExternalReaderContext(extfs),
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := readDeltaRecords(reader); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if len(legacyPaths) > 0 {
|
|
reader, err := storage.NewDeltalogReader(
|
|
pkField.DataType,
|
|
legacyPaths,
|
|
storage.WithVersion(storage.StorageV1),
|
|
storage.WithDownloader(func(ctx context.Context, paths []string) ([][]byte, error) {
|
|
return readExternalFiles(ctx, createStorageConfig(), extfs, paths)
|
|
}),
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := readDeltaRecords(reader); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
// V3: delta data lives in manifest.
|
|
paths, err := packed.GetDeltaLogPathsFromManifest(manifestPath, createStorageConfig())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := readPaths(paths,
|
|
storage.WithStorageConfig(createStorageConfig()),
|
|
storage.WithVersion(storage.StorageV3),
|
|
); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
} else {
|
|
// V1: delta data referenced by Deltalogs entries
|
|
var paths []string
|
|
for _, deltalog := range deltaLogs {
|
|
for _, binlog := range lo.Filter(deltalog.Binlogs, valid) {
|
|
if p := binlog.GetLogPath(); p != "" {
|
|
paths = append(paths, p)
|
|
}
|
|
}
|
|
}
|
|
if err := readPaths(paths,
|
|
storage.WithDownloader(func(ctx context.Context, paths []string) ([][]byte, error) {
|
|
return loader.cm.MultiRead(ctx, paths)
|
|
}),
|
|
); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
err = segment.LoadDeltaData(ctx, deltaData)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
mlog.Info(context.TODO(), "load delta logs done", mlog.Int64("deleteCount", deltaData.DeleteRowCount()))
|
|
return nil
|
|
}
|
|
|
|
func milvusTableDeltalogPaths(deltaLogs []*datapb.FieldBinlog) []string {
|
|
paths := make([]string, 0)
|
|
for _, deltaLog := range deltaLogs {
|
|
for _, binlog := range deltaLog.GetBinlogs() {
|
|
if binlog.GetLogPath() != "" {
|
|
paths = append(paths, binlog.GetLogPath())
|
|
}
|
|
}
|
|
}
|
|
return paths
|
|
}
|
|
|
|
// validateMilvusTableRealPKDeltalogPaths rejects target-owned deltalogs from a
|
|
// real-PK milvus-table manifest. Real-PK load may consume source StorageV3
|
|
// deltas and legacy snapshot L0 deltas; target-owned deltas are reserved for
|
|
// virtual-PK translation.
|
|
func validateMilvusTableRealPKDeltalogPaths(manifestPath string, deltaPaths []string) error {
|
|
basePath, _, err := packed.UnmarshalManifestPath(manifestPath)
|
|
if err != nil {
|
|
return merr.WrapErrServiceInternalErr(err, "parse milvus-table manifest path")
|
|
}
|
|
targetDeltaPrefix := strings.TrimRight(basePath, "/") + "/_delta/"
|
|
for _, deltaPath := range deltaPaths {
|
|
if deltaPath == "" {
|
|
continue
|
|
}
|
|
if strings.HasPrefix(deltaPath, targetDeltaPrefix) {
|
|
return merr.WrapErrServiceInternalMsg("milvus-table real-PK manifest must not contain target-owned deltalog %s", deltaPath)
|
|
}
|
|
if err := packed.ValidateMilvusTableSourceDeltalogPath(deltaPath); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// readExternalFiles reads whole files through packed external-spec filesystem
|
|
// aliases and checks ctx before each potentially large read.
|
|
func readExternalFiles(
|
|
ctx context.Context,
|
|
storageConfig *indexpb.StorageConfig,
|
|
extfs packed.ExternalSpecContext,
|
|
paths []string,
|
|
) ([][]byte, error) {
|
|
data := make([][]byte, len(paths))
|
|
for i, path := range paths {
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
content, err := packed.ReadFileWithExternalSpec(storageConfig, path, extfs)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
data[i] = content
|
|
}
|
|
return data, nil
|
|
}
|
|
|
|
// LoadDeltaLogs load deltalog and write delta data into provided segment.
|
|
// it also executes resource protection logic in case of OOM.
|
|
func (loader *segmentLoader) LoadDeltaLogs(ctx context.Context, segment Segment, loadInfo *querypb.SegmentLoadInfo) error {
|
|
// Check memory & storage limit
|
|
requestResourceResult, err := loader.requestResource(ctx, loadInfo)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "request resource failed", mlog.Err(err))
|
|
return err
|
|
}
|
|
defer loader.freeRequestResource(requestResourceResult)
|
|
return loader.loadDeltalogs(ctx, segment, loadInfo)
|
|
}
|
|
|
|
func createStorageConfig() *indexpb.StorageConfig {
|
|
params := paramtable.Get()
|
|
if params.CommonCfg.StorageType.GetValue() == "local" {
|
|
return &indexpb.StorageConfig{
|
|
RootPath: params.LocalStorageCfg.Path.GetValue(),
|
|
StorageType: params.CommonCfg.StorageType.GetValue(),
|
|
}
|
|
}
|
|
return &indexpb.StorageConfig{
|
|
Address: params.MinioCfg.Address.GetValue(),
|
|
AccessKeyID: params.MinioCfg.AccessKeyID.GetValue(),
|
|
SecretAccessKey: params.MinioCfg.SecretAccessKey.GetValue(),
|
|
UseSSL: params.MinioCfg.UseSSL.GetAsBool(),
|
|
SslCACert: params.MinioCfg.SslCACert.GetValue(),
|
|
BucketName: params.MinioCfg.BucketName.GetValue(),
|
|
RootPath: params.MinioCfg.RootPath.GetValue(),
|
|
UseIAM: params.MinioCfg.UseIAM.GetAsBool(),
|
|
IAMEndpoint: params.MinioCfg.IAMEndpoint.GetValue(),
|
|
StorageType: params.CommonCfg.StorageType.GetValue(),
|
|
Region: params.MinioCfg.Region.GetValue(),
|
|
UseVirtualHost: params.MinioCfg.UseVirtualHost.GetAsBool(),
|
|
CloudProvider: params.MinioCfg.CloudProvider.GetValue(),
|
|
RequestTimeoutMs: params.MinioCfg.RequestTimeoutMs.GetAsInt64(),
|
|
GcpCredentialJSON: params.MinioCfg.GcpCredentialJSON.GetValue(),
|
|
SslTlsMinVersion: params.MinioCfg.SslTLSMinVersion.GetValue(),
|
|
UseCrc32CChecksum: params.MinioCfg.UseCRC32C.GetAsBool(),
|
|
}
|
|
}
|
|
|
|
func (loader *segmentLoader) patchEntryNumber(ctx context.Context, segment *LocalSegment, loadInfo *querypb.SegmentLoadInfo) error {
|
|
var needReset bool
|
|
|
|
segment.fieldIndexes.Range(func(indexID int64, info *IndexedFieldInfo) bool {
|
|
for _, info := range info.FieldBinlog.GetBinlogs() {
|
|
if info.GetEntriesNum() == 0 {
|
|
needReset = true
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
})
|
|
if !needReset {
|
|
return nil
|
|
}
|
|
|
|
mlog.Warn(context.TODO(), "legacy segment binlog found, start to patch entry num", mlog.Int64("segmentID", segment.ID()))
|
|
rowIDField := lo.FindOrElse(loadInfo.BinlogPaths, nil, func(binlog *datapb.FieldBinlog) bool {
|
|
return binlog.GetFieldID() == common.RowIDField
|
|
})
|
|
|
|
if rowIDField == nil {
|
|
return merr.WrapErrDataIntegrityMsg("rowID field binlog not found")
|
|
}
|
|
|
|
counts := make([]int64, 0, len(rowIDField.GetBinlogs()))
|
|
for _, binlog := range rowIDField.GetBinlogs() {
|
|
// binlog.LogPath has already been filled
|
|
bs, err := loader.cm.Read(ctx, binlog.LogPath)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// get binlog entry num from rowID field
|
|
// since header does not store entry numb, we have to read all data here
|
|
|
|
reader, err := storage.NewBinlogReader(bs)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
er, err := reader.NextEventReader()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
rowIDs, _, err := er.GetInt64FromPayload()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
counts = append(counts, int64(len(rowIDs)))
|
|
}
|
|
|
|
var err error
|
|
segment.fieldIndexes.Range(func(indexID int64, info *IndexedFieldInfo) bool {
|
|
if len(info.FieldBinlog.GetBinlogs()) != len(counts) {
|
|
err = merr.WrapErrDataIntegrityMsg("rowID & index binlog number not matched")
|
|
return false
|
|
}
|
|
for i, binlog := range info.FieldBinlog.GetBinlogs() {
|
|
binlog.EntriesNum = counts[i]
|
|
}
|
|
return true
|
|
})
|
|
return err
|
|
}
|
|
|
|
// JoinIDPath joins ids to path format.
|
|
func JoinIDPath(ids ...int64) string {
|
|
idStr := make([]string, 0, len(ids))
|
|
for _, id := range ids {
|
|
idStr = append(idStr, strconv.FormatInt(id, 10))
|
|
}
|
|
return path.Join(idStr...)
|
|
}
|
|
|
|
// After introducing the caching layer's lazy loading and eviction mechanisms, most parts of a segment won't be
|
|
// loaded into memory or disk immediately, even if the segment is marked as LOADED. This means physical resource
|
|
// usage may be very low.
|
|
// However, we still need to reserve enough resources for the segments marked as LOADED. The reserved resource is
|
|
// treated as the logical resource usage. Logical resource usage is based on the segment final resource usage.
|
|
// checkLogicalSegmentSize checks whether the memory & disk is sufficient to load the segments,
|
|
// returns the memory & disk logical usage while loading if possible to load, otherwise, returns error
|
|
func (loader *segmentLoader) checkLogicalSegmentSize(ctx context.Context, segmentLoadInfos []*querypb.SegmentLoadInfo, totalMem uint64) (uint64, uint64, error) {
|
|
if !paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool() {
|
|
return 0, 0, nil
|
|
}
|
|
|
|
if len(segmentLoadInfos) == 0 {
|
|
return 0, 0, nil
|
|
}
|
|
|
|
logicalMemUsage := loader.manager.Segment.GetLogicalResource().MemorySize
|
|
logicalDiskUsage := loader.manager.Segment.GetLogicalResource().DiskSize
|
|
|
|
logicalMemUsage += loader.committedLogicalResource.MemorySize
|
|
logicalDiskUsage += loader.committedLogicalResource.DiskSize
|
|
|
|
// logical resource usage is based on the segment final resource usage,
|
|
// so we need to estimate the final resource usage of the segments
|
|
finalFactor := resourceEstimateFactor{
|
|
deltaDataExpansionFactor: paramtable.Get().QueryNodeCfg.DeltaDataExpansionRate.GetAsFloat(),
|
|
textIndexExpansionFactor: paramtable.Get().QueryNodeCfg.TextIndexExpansionFactor.GetAsFloat(),
|
|
TieredEvictionEnabled: paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool(),
|
|
TieredEvictableMemoryCacheRatio: paramtable.Get().QueryNodeCfg.TieredEvictableMemoryCacheRatio.GetAsFloat(),
|
|
TieredEvictableDiskCacheRatio: paramtable.Get().QueryNodeCfg.TieredEvictableDiskCacheRatio.GetAsFloat(),
|
|
}
|
|
predictLogicalMemUsage := logicalMemUsage
|
|
predictLogicalDiskUsage := logicalDiskUsage
|
|
for _, loadInfo := range segmentLoadInfos {
|
|
collection := loader.manager.Collection.Get(loadInfo.GetCollectionID())
|
|
finalUsage, err := estimateLogicalResourceUsageOfSegment(collection.Schema(), loadInfo, finalFactor)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "failed to estimate final resource usage of segment",
|
|
mlog.Int64("collectionID", loadInfo.GetCollectionID()),
|
|
mlog.Int64("segmentID", loadInfo.GetSegmentID()),
|
|
mlog.Err(err))
|
|
return 0, 0, err
|
|
}
|
|
|
|
mlog.Debug(context.TODO(), "segment logical resource for loading",
|
|
mlog.Int64("segmentID", loadInfo.GetSegmentID()),
|
|
mlog.Float64("memoryUsage(MB)", logutil.ToMB(float64(finalUsage.MemorySize))),
|
|
mlog.Float64("diskUsage(MB)", logutil.ToMB(float64(finalUsage.DiskSize))),
|
|
)
|
|
predictLogicalDiskUsage += finalUsage.DiskSize
|
|
predictLogicalMemUsage += finalUsage.MemorySize
|
|
}
|
|
|
|
mlog.Info(context.TODO(), "predict memory and disk logical usage after loaded (in MiB)",
|
|
mlog.Float64("predictLogicalMemUsage(MB)", logutil.ToMB(float64(predictLogicalMemUsage))),
|
|
mlog.Float64("predictLogicalDiskUsage(MB)", logutil.ToMB(float64(predictLogicalDiskUsage))),
|
|
)
|
|
|
|
logicalMemUsageLimit := uint64(float64(totalMem) * paramtable.Get().QueryNodeCfg.OverloadedMemoryThresholdPercentage.GetAsFloat())
|
|
logicalDiskUsageLimit := uint64(float64(paramtable.Get().QueryNodeCfg.DiskCapacityLimit.GetAsInt64()) * paramtable.Get().QueryNodeCfg.MaxDiskUsagePercentage.GetAsFloat())
|
|
|
|
if predictLogicalMemUsage > logicalMemUsageLimit {
|
|
mlog.Warn(context.TODO(), "logical memory usage checking for segment loading failed",
|
|
mlog.String("resourceType", "Memory"),
|
|
mlog.Float64("predictLogicalMemUsageMB", logutil.ToMB(float64(predictLogicalMemUsage))),
|
|
mlog.Float64("logicalMemUsageLimitMB", logutil.ToMB(float64(logicalMemUsageLimit))),
|
|
mlog.Float64("evictableMemoryCacheRatio", paramtable.Get().QueryNodeCfg.TieredEvictableMemoryCacheRatio.GetAsFloat()),
|
|
)
|
|
return 0, 0, merr.WrapErrSegmentRequestResourceFailed("Memory")
|
|
}
|
|
|
|
if predictLogicalDiskUsage > logicalDiskUsageLimit {
|
|
mlog.Warn(ctx, fmt.Sprintf("Logical disk usage checking for segment loading failed, predictLogicalDiskUsage = %v MB, LogicalDiskUsageLimit = %v MB, decrease the evictableDiskCacheRatio (current: %v) if you want to load more segments",
|
|
logutil.ToMB(float64(predictLogicalDiskUsage)),
|
|
logutil.ToMB(float64(logicalDiskUsageLimit)),
|
|
paramtable.Get().QueryNodeCfg.TieredEvictableDiskCacheRatio.GetAsFloat(),
|
|
))
|
|
return 0, 0, merr.WrapErrSegmentRequestResourceFailed("Disk")
|
|
}
|
|
|
|
return predictLogicalMemUsage - logicalMemUsage, predictLogicalDiskUsage - logicalDiskUsage, nil
|
|
}
|
|
|
|
func (loader *segmentLoader) estimateSegmentLoadingResourceUsage(ctx context.Context, segmentLoadInfos ...*querypb.SegmentLoadInfo) (*ResourceUsage, uint64, error) {
|
|
if len(segmentLoadInfos) == 0 {
|
|
return &ResourceUsage{}, 0, nil
|
|
}
|
|
|
|
logger := mlog.With(
|
|
mlog.Int64("collectionID", segmentLoadInfos[0].GetCollectionID()),
|
|
)
|
|
|
|
maxFactor := resourceEstimateFactor{
|
|
memoryUsageFactor: paramtable.Get().QueryNodeCfg.LoadMemoryUsageFactor.GetAsFloat(),
|
|
memoryIndexUsageFactor: paramtable.Get().QueryNodeCfg.MemoryIndexLoadPredictMemoryUsageFactor.GetAsFloat(),
|
|
EnableInterminSegmentIndex: paramtable.Get().QueryNodeCfg.EnableInterminSegmentIndex.GetAsBool(),
|
|
tempSegmentIndexFactor: paramtable.Get().QueryNodeCfg.InterimIndexMemExpandRate.GetAsFloat(),
|
|
deltaDataExpansionFactor: paramtable.Get().QueryNodeCfg.DeltaDataExpansionRate.GetAsFloat(),
|
|
jsonKeyStatsExpansionFactor: paramtable.Get().QueryNodeCfg.JSONKeyStatsExpansionFactor.GetAsFloat(),
|
|
textIndexExpansionFactor: paramtable.Get().QueryNodeCfg.TextIndexExpansionFactor.GetAsFloat(),
|
|
TieredEvictionEnabled: paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool(),
|
|
externalRawDataFactor: paramtable.Get().QueryNodeCfg.ExternalCollectionRawDataFactor.GetAsFloat(),
|
|
}
|
|
maxSegmentSize := uint64(0)
|
|
predictMemUsage := uint64(0)
|
|
predictDiskUsage := uint64(0)
|
|
var predictGpuMemUsage []uint64
|
|
mmapFieldCount := 0
|
|
for _, loadInfo := range segmentLoadInfos {
|
|
collection := loader.manager.Collection.Get(loadInfo.GetCollectionID())
|
|
loadingUsage, err := estimateLoadingResourceUsageOfSegment(collection.Schema(), loadInfo, maxFactor)
|
|
if err != nil {
|
|
logger.Warn(ctx, "failed to estimate max resource usage of segment",
|
|
mlog.Int64("collectionID", loadInfo.GetCollectionID()),
|
|
mlog.Int64("segmentID", loadInfo.GetSegmentID()),
|
|
mlog.Err(err))
|
|
return nil, 0, err
|
|
}
|
|
|
|
logger.Debug(ctx, "segment resource for loading",
|
|
mlog.Int64("segmentID", loadInfo.GetSegmentID()),
|
|
mlog.Float64("loadingMemoryUsage(MB)", logutil.ToMB(float64(loadingUsage.MemorySize))),
|
|
mlog.Float64("loadingDiskUsage(MB)", logutil.ToMB(float64(loadingUsage.DiskSize))),
|
|
mlog.Float64("memoryLoadFactor", maxFactor.memoryUsageFactor),
|
|
)
|
|
mmapFieldCount += loadingUsage.MmapFieldCount
|
|
predictDiskUsage += loadingUsage.DiskSize
|
|
predictMemUsage += loadingUsage.MemorySize
|
|
predictGpuMemUsage = append(predictGpuMemUsage, loadingUsage.FieldGpuMemorySize...)
|
|
if loadingUsage.MemorySize > maxSegmentSize {
|
|
maxSegmentSize = loadingUsage.MemorySize
|
|
}
|
|
}
|
|
|
|
return &ResourceUsage{
|
|
MemorySize: predictMemUsage,
|
|
DiskSize: predictDiskUsage,
|
|
MmapFieldCount: mmapFieldCount,
|
|
FieldGpuMemorySize: predictGpuMemUsage,
|
|
}, maxSegmentSize, nil
|
|
}
|
|
|
|
// checkLoadingResource checks physical resource limits for an already-estimated loading usage.
|
|
// Callers that race with load resource commits must hold loader.mut.
|
|
func (loader *segmentLoader) checkLoadingResource(
|
|
ctx context.Context,
|
|
logger *mlog.Logger,
|
|
loadingUsage *ResourceUsage,
|
|
maxSegmentSize uint64,
|
|
totalMem uint64,
|
|
memUsage uint64,
|
|
localDiskUsage int64,
|
|
) error {
|
|
memUsage += loader.committedResource.MemorySize
|
|
if memUsage == 0 || totalMem == 0 {
|
|
return merr.WrapErrServiceInternalMsg("get memory failed when checkLoadingResource")
|
|
}
|
|
|
|
diskUsage := uint64(localDiskUsage) + loader.committedResource.DiskSize
|
|
predictMemUsage := memUsage + loadingUsage.MemorySize
|
|
predictDiskUsage := diskUsage + loadingUsage.DiskSize
|
|
|
|
logger.Info(ctx, "predict memory and disk usage while loading (in MiB)",
|
|
mlog.Float64("maxSegmentSize(MB)", logutil.ToMB(float64(maxSegmentSize))),
|
|
mlog.Float64("committedMemSize(MB)", logutil.ToMB(float64(loader.committedResource.MemorySize))),
|
|
mlog.Float64("memLimit(MB)", logutil.ToMB(float64(totalMem))),
|
|
mlog.Float64("memUsage(MB)", logutil.ToMB(float64(memUsage))),
|
|
mlog.Float64("committedDiskSize(MB)", logutil.ToMB(float64(loader.committedResource.DiskSize))),
|
|
mlog.Float64("diskUsage(MB)", logutil.ToMB(float64(diskUsage))),
|
|
mlog.Float64("predictMemUsage(MB)", logutil.ToMB(float64(predictMemUsage))),
|
|
mlog.Float64("predictDiskUsage(MB)", logutil.ToMB(float64(predictDiskUsage))),
|
|
mlog.Int("mmapFieldCount", loadingUsage.MmapFieldCount),
|
|
)
|
|
|
|
var loadingResource C.CResourceUsage
|
|
reservedLoadingResource := false
|
|
if paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool() {
|
|
loadingResource = C.CResourceUsage{
|
|
memory_bytes: C.int64_t(loadingUsage.MemorySize),
|
|
disk_bytes: C.int64_t(loadingUsage.DiskSize),
|
|
}
|
|
|
|
// try to reserve loading resource from caching layer
|
|
if ok := C.TryReserveLoadingResourceWithTimeout(loadingResource, 1000); !ok {
|
|
return merr.WrapErrSegmentRequestResourceFailed("memory/disk",
|
|
fmt.Sprintf("failed to reserve loading resource from caching layer, predictMemUsage = %v MB, predictDiskUsage = %v MB, memUsage = %v MB, diskUsage = %v MB, memoryThresholdFactor = %f, diskThresholdFactor = %f",
|
|
logutil.ToMB(float64(predictMemUsage)),
|
|
logutil.ToMB(float64(predictDiskUsage)),
|
|
logutil.ToMB(float64(memUsage)),
|
|
logutil.ToMB(float64(diskUsage)),
|
|
paramtable.Get().QueryNodeCfg.OverloadedMemoryThresholdPercentage.GetAsFloat(),
|
|
paramtable.Get().QueryNodeCfg.MaxDiskUsagePercentage.GetAsFloat(),
|
|
))
|
|
}
|
|
reservedLoadingResource = true
|
|
} else {
|
|
// fallback to original segment loading logic
|
|
if predictMemUsage > uint64(float64(totalMem)*paramtable.Get().QueryNodeCfg.OverloadedMemoryThresholdPercentage.GetAsFloat()) {
|
|
mlog.Warn(context.TODO(), "load segment failed, OOM if load",
|
|
mlog.String("resourceType", "Memory"),
|
|
mlog.Float64("maxSegmentSizeMB", logutil.ToMB(float64(maxSegmentSize))),
|
|
mlog.Float64("memUsageMB", logutil.ToMB(float64(memUsage))),
|
|
mlog.Float64("predictMemUsageMB", logutil.ToMB(float64(predictMemUsage))),
|
|
mlog.Float64("totalMemMB", logutil.ToMB(float64(totalMem))),
|
|
mlog.Float64("thresholdFactor", paramtable.Get().QueryNodeCfg.OverloadedMemoryThresholdPercentage.GetAsFloat()),
|
|
)
|
|
return merr.WrapErrSegmentRequestResourceFailed("Memory")
|
|
}
|
|
|
|
if predictDiskUsage > uint64(float64(paramtable.Get().QueryNodeCfg.DiskCapacityLimit.GetAsInt64())*paramtable.Get().QueryNodeCfg.MaxDiskUsagePercentage.GetAsFloat()) {
|
|
mlog.Warn(context.TODO(), "load segment failed, disk space is not enough",
|
|
mlog.String("resourceType", "Disk"),
|
|
mlog.Float64("diskUsageMB", logutil.ToMB(float64(diskUsage))),
|
|
mlog.Float64("predictDiskUsageMB", logutil.ToMB(float64(predictDiskUsage))),
|
|
mlog.Float64("totalDiskMB", logutil.ToMB(float64(uint64(paramtable.Get().QueryNodeCfg.DiskCapacityLimit.GetAsInt64())))),
|
|
mlog.Float64("thresholdFactor", paramtable.Get().QueryNodeCfg.MaxDiskUsagePercentage.GetAsFloat()),
|
|
)
|
|
return merr.WrapErrSegmentRequestResourceFailed("Disk")
|
|
}
|
|
}
|
|
|
|
err := checkSegmentGpuMemSize(loadingUsage.FieldGpuMemorySize, float32(paramtable.Get().GpuConfig.OverloadedMemoryThresholdPercentage.GetAsFloat()))
|
|
if err != nil {
|
|
if reservedLoadingResource {
|
|
C.ReleaseLoadingResource(loadingResource)
|
|
}
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// this function is used to estimate the logical resource usage of a segment, which should only be used when tiered eviction is enabled
|
|
// the result is the final resource usage of the segment inevictable part plus the final usage of evictable part with cache ratio applied
|
|
// TODO: the inevictable part is not correct, since we cannot know the final resource usage of interim index and default-value column before loading,
|
|
// current they are ignored, but we should consider them in the future
|
|
func estimateLogicalResourceUsageOfSegment(schema *schemapb.CollectionSchema, loadInfo *querypb.SegmentLoadInfo, multiplyFactor resourceEstimateFactor) (usage *ResourceUsage, err error) {
|
|
var segmentInevictableMemorySize, segmentInevictableDiskSize uint64
|
|
var segmentEvictableMemorySize, segmentEvictableDiskSize uint64
|
|
|
|
id2Binlogs := lo.SliceToMap(loadInfo.BinlogPaths, func(fieldBinlog *datapb.FieldBinlog) (int64, *datapb.FieldBinlog) {
|
|
return fieldBinlog.GetFieldID(), fieldBinlog
|
|
})
|
|
|
|
schemaHelper, err := typeutil.CreateSchemaHelper(schema)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "failed to create schema helper", mlog.String("name", schema.GetName()), mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
ctx := context.TODO()
|
|
|
|
// PART 1: calculate logical resource usage of indexes
|
|
for _, fieldIndexInfo := range loadInfo.IndexInfos {
|
|
fieldID := fieldIndexInfo.GetFieldID()
|
|
if len(fieldIndexInfo.GetIndexFilePaths()) > 0 {
|
|
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
isVectorType := typeutil.IsVectorType(fieldSchema.GetDataType())
|
|
|
|
var estimateResult ResourceEstimate
|
|
err = GetCLoadInfoWithFunc(ctx, fieldSchema, loadInfo, fieldIndexInfo, func(c *LoadIndexInfo) error {
|
|
GetDynamicPool().Submit(func() (any, error) {
|
|
loadResourceRequest := C.EstimateLoadIndexResource(c.cLoadIndexInfo)
|
|
estimateResult = GetResourceEstimate(&loadResourceRequest)
|
|
return nil, nil
|
|
}).Await()
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, merr.Wrapf(err, "failed to estimate logical resource usage of index, collection %d, segment %d, indexBuildID %d",
|
|
loadInfo.GetCollectionID(),
|
|
loadInfo.GetSegmentID(),
|
|
fieldIndexInfo.GetBuildID())
|
|
}
|
|
segmentEvictableMemorySize += estimateResult.FinalMemoryCost
|
|
segmentEvictableDiskSize += estimateResult.FinalDiskCost
|
|
|
|
// could skip binlog or
|
|
// could be missing for new field or storage v2 group 0
|
|
if estimateResult.HasRawData &&
|
|
!paramtable.Get().QueryNodeCfg.PreferFieldDataWhenIndexHasRawData.GetAsBool() {
|
|
delete(id2Binlogs, fieldID)
|
|
continue
|
|
}
|
|
|
|
// BM25 only checks vector datatype
|
|
// scalar index does not have metrics type key
|
|
if !isVectorType {
|
|
continue
|
|
}
|
|
|
|
metricType, err := funcutil.GetAttrByKeyFromRepeatedKV(common.MetricTypeKey, fieldIndexInfo.IndexParams)
|
|
if err != nil {
|
|
return nil, merr.Wrapf(err, "failed to estimate logical resource usage of index, metric type not found, collection %d, segment %d, indexBuildID %d",
|
|
loadInfo.GetCollectionID(),
|
|
loadInfo.GetSegmentID(),
|
|
fieldIndexInfo.GetBuildID())
|
|
}
|
|
// skip raw data for BM25 index
|
|
if metricType == metric.BM25 {
|
|
delete(id2Binlogs, fieldID)
|
|
}
|
|
}
|
|
}
|
|
|
|
// PART 2: calculate logical resource usage of binlogs
|
|
for fieldID, fieldBinlog := range id2Binlogs {
|
|
fieldIDs := fieldBinlog.GetChildFields()
|
|
// legacy default split
|
|
if len(fieldIDs) == 0 {
|
|
fieldIDs = []int64{fieldID}
|
|
}
|
|
binlogSize := uint64(getBinlogDataMemorySize(fieldBinlog))
|
|
|
|
var supportInterimIndexDataType bool
|
|
var containsTimestampField bool
|
|
var doubleMemoryDataField bool
|
|
var legacyNilSchema bool
|
|
mmapEnabled := true
|
|
isVectorType := true
|
|
|
|
for _, fieldID := range fieldIDs {
|
|
// get field schema from fieldID
|
|
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "failed to get field schema", mlog.Int64("fieldID", fieldID), mlog.String("name", schema.GetName()), mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
|
|
// missing mapping, shall be "0" group for storage v2
|
|
if fieldSchema == nil {
|
|
legacyNilSchema = true
|
|
break
|
|
}
|
|
|
|
supportInterimIndexDataType = supportInterimIndexDataType || SupportInterimIndexDataType(fieldSchema.GetDataType())
|
|
isVectorType = isVectorType && typeutil.IsVectorType(fieldSchema.GetDataType())
|
|
// constainSystemField = constainSystemField || common.IsSystemField(fieldSchema.GetFieldID())
|
|
mmapEnabled = mmapEnabled && isDataMmapEnable(fieldSchema)
|
|
containsTimestampField = containsTimestampField || DoubleMemorySystemField(fieldSchema.GetFieldID())
|
|
doubleMemoryDataField = doubleMemoryDataField || DoubleMemoryDataType(fieldSchema.GetDataType())
|
|
}
|
|
|
|
// TODO: add default-value column's resource usage to inevictable part
|
|
// TODO: add interim index's resource usage to inevictable part
|
|
|
|
if legacyNilSchema {
|
|
segmentEvictableMemorySize += binlogSize
|
|
continue
|
|
}
|
|
|
|
// timestamp field double in InsertRecord & TimestampIndex
|
|
if containsTimestampField {
|
|
timestampSize := lo.SumBy(fieldBinlog.GetBinlogs(), func(binlog *datapb.Binlog) int64 {
|
|
return binlog.GetEntriesNum() * 4
|
|
})
|
|
segmentInevictableMemorySize += 2 * uint64(timestampSize)
|
|
}
|
|
|
|
if isVectorType {
|
|
mmapVectorField := paramtable.Get().QueryNodeCfg.MmapVectorField.GetAsBool()
|
|
if mmapVectorField {
|
|
segmentEvictableDiskSize += binlogSize
|
|
} else {
|
|
segmentEvictableMemorySize += binlogSize
|
|
}
|
|
} else if !mmapEnabled {
|
|
segmentEvictableMemorySize += binlogSize
|
|
if doubleMemoryDataField {
|
|
segmentEvictableMemorySize += binlogSize
|
|
}
|
|
} else {
|
|
segmentEvictableDiskSize += binlogSize
|
|
}
|
|
}
|
|
|
|
// PART 3: calculate logical resource usage of stats data
|
|
for _, fieldBinlog := range loadInfo.Statslogs {
|
|
segmentInevictableMemorySize += uint64(getBinlogDataMemorySize(fieldBinlog))
|
|
}
|
|
|
|
// PART 4: calculate logical resource usage of delete data
|
|
for _, fieldBinlog := range loadInfo.Deltalogs {
|
|
// MemorySize of filedBinlog is the actual size in memory, so the expansionFactor
|
|
// should be 1, in most cases.
|
|
expansionFactor := float64(1)
|
|
memSize := getBinlogDataMemorySize(fieldBinlog)
|
|
|
|
// Note: If MemorySize == DiskSize, it means the segment comes from Milvus 2.3,
|
|
// MemorySize is actually compressed DiskSize of deltalog, so we'll fallback to use
|
|
// deltaExpansionFactor to compromise the compression ratio.
|
|
if memSize == getBinlogDataDiskSize(fieldBinlog) {
|
|
expansionFactor = multiplyFactor.deltaDataExpansionFactor
|
|
}
|
|
segmentInevictableMemorySize += uint64(float64(memSize) * expansionFactor)
|
|
}
|
|
|
|
// PART 5: calculate logical resource usage of text index stats data
|
|
// Text match indexes are evictable (support_eviction=true in caching layer).
|
|
// Text match index mmap is driven by scalar_field_enable_mmap (same as raw scalar data).
|
|
textIndexMmapEnable := paramtable.Get().QueryNodeCfg.MmapScalarField.GetAsBool()
|
|
for _, textStats := range loadInfo.GetTextStatsLogs() {
|
|
if textIndexMmapEnable {
|
|
segmentEvictableDiskSize += uint64(float64(textStats.GetMemorySize()) * multiplyFactor.textIndexExpansionFactor)
|
|
} else {
|
|
segmentEvictableMemorySize += uint64(float64(textStats.GetMemorySize()) * multiplyFactor.textIndexExpansionFactor)
|
|
}
|
|
}
|
|
|
|
mlog.Debug(context.TODO(), "estimate logical resoure usage result",
|
|
mlog.Int64("segmentID", loadInfo.GetSegmentID()),
|
|
mlog.Uint64("segmentInevictableMemorySize", segmentInevictableMemorySize),
|
|
mlog.Uint64("segmentEvictableMemorySize", segmentEvictableMemorySize),
|
|
mlog.Uint64("segmentInevictableDiskSize", segmentInevictableDiskSize),
|
|
mlog.Uint64("segmentEvictableDiskSize", segmentEvictableDiskSize),
|
|
)
|
|
|
|
return &ResourceUsage{
|
|
MemorySize: segmentInevictableMemorySize + uint64(float64(segmentEvictableMemorySize)*multiplyFactor.TieredEvictableMemoryCacheRatio),
|
|
DiskSize: segmentInevictableDiskSize + uint64(float64(segmentEvictableDiskSize)*multiplyFactor.TieredEvictableDiskCacheRatio),
|
|
}, nil
|
|
}
|
|
|
|
// estimateLoadingResourceUsageOfSegment estimates the resource usage of the segment when loading,
|
|
// it will return two different results, depending on the value of tiered eviction parameter:
|
|
// - when tiered eviction is enabled, the result is the max resource usage of the segment that cannot be managed by caching layer,
|
|
// which should be a subset of the segment inevictable part
|
|
// - when tiered eviction is disabled, the result is the max resource usage of both the segment evictable and inevictable part
|
|
func estimateLoadingResourceUsageOfSegment(schema *schemapb.CollectionSchema, loadInfo *querypb.SegmentLoadInfo, multiplyFactor resourceEstimateFactor) (usage *ResourceUsage, err error) {
|
|
var segMemoryLoadingSize, segDiskLoadingSize uint64
|
|
var indexMemorySize uint64
|
|
var mmapFieldCount int
|
|
var fieldGpuMemorySize []uint64
|
|
|
|
id2Binlogs := lo.SliceToMap(loadInfo.BinlogPaths, func(fieldBinlog *datapb.FieldBinlog) (int64, *datapb.FieldBinlog) {
|
|
return fieldBinlog.GetFieldID(), fieldBinlog
|
|
})
|
|
|
|
schemaHelper, err := typeutil.CreateSchemaHelper(schema)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "failed to create schema helper", mlog.String("name", schema.GetName()), mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
indexedFields := make(map[int64]struct{})
|
|
ctx := context.Background()
|
|
|
|
// PART 1: calculate size of indexes
|
|
for _, fieldIndexInfo := range loadInfo.IndexInfos {
|
|
fieldID := fieldIndexInfo.GetFieldID()
|
|
if len(fieldIndexInfo.GetIndexFilePaths()) > 0 {
|
|
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
|
|
if err != nil {
|
|
// field might have been dropped, skip its index
|
|
mlog.Info(ctx, "skip index for dropped field", mlog.FieldFieldID(fieldID), mlog.String("name", schema.GetName()))
|
|
continue
|
|
}
|
|
indexedFields[fieldID] = struct{}{}
|
|
|
|
isVectorType := typeutil.IsVectorType(fieldSchema.GetDataType())
|
|
|
|
var estimateResult ResourceEstimate
|
|
err = GetCLoadInfoWithFunc(ctx, fieldSchema, loadInfo, fieldIndexInfo, func(c *LoadIndexInfo) error {
|
|
GetDynamicPool().Submit(func() (any, error) {
|
|
loadResourceRequest := C.EstimateLoadIndexResource(c.cLoadIndexInfo)
|
|
estimateResult = GetResourceEstimate(&loadResourceRequest)
|
|
return nil, nil
|
|
}).Await()
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, merr.Wrapf(err, "failed to estimate loading resource usage of index, collection %d, segment %d, indexBuildID %d",
|
|
loadInfo.GetCollectionID(),
|
|
loadInfo.GetSegmentID(),
|
|
fieldIndexInfo.GetBuildID())
|
|
}
|
|
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
indexMemorySize += estimateResult.MaxMemoryCost
|
|
segDiskLoadingSize += estimateResult.MaxDiskCost
|
|
}
|
|
|
|
if gpuIndexRequiresGpu(fieldIndexInfo.IndexParams) {
|
|
fieldGpuMemorySize = append(fieldGpuMemorySize, estimateResult.MaxMemoryCost)
|
|
}
|
|
|
|
// could skip binlog or
|
|
// could be missing for new field or storage v2 group 0
|
|
if estimateResult.HasRawData &&
|
|
!paramtable.Get().QueryNodeCfg.PreferFieldDataWhenIndexHasRawData.GetAsBool() {
|
|
delete(id2Binlogs, fieldID)
|
|
continue
|
|
}
|
|
|
|
// BM25 only checks vector datatype
|
|
// scalar index does not have metrics type key
|
|
if !isVectorType {
|
|
continue
|
|
}
|
|
|
|
metricType, err := funcutil.GetAttrByKeyFromRepeatedKV(common.MetricTypeKey, fieldIndexInfo.IndexParams)
|
|
if err != nil {
|
|
return nil, merr.Wrapf(err, "failed to estimate loading resource usage of index, metric type not found, collection %d, segment %d, indexBuildID %d",
|
|
loadInfo.GetCollectionID(),
|
|
loadInfo.GetSegmentID(),
|
|
fieldIndexInfo.GetBuildID())
|
|
}
|
|
// skip raw data for BM25 index
|
|
if metricType == metric.BM25 {
|
|
delete(id2Binlogs, fieldID)
|
|
}
|
|
}
|
|
}
|
|
|
|
// PART 2: calculate size of binlogs
|
|
for fieldID, fieldBinlog := range id2Binlogs {
|
|
fieldIDs := fieldBinlog.GetChildFields()
|
|
// legacy default split
|
|
if len(fieldIDs) == 0 {
|
|
fieldIDs = []int64{fieldID}
|
|
}
|
|
binlogSize := uint64(getBinlogDataMemorySize(fieldBinlog))
|
|
|
|
var supportInterimIndexDataType bool
|
|
var containsTimestampField bool
|
|
var doubleMomoryDataField bool
|
|
var legacyNilSchema bool
|
|
mmapEnabled := true
|
|
isVectorType := true
|
|
hasIndex := true
|
|
|
|
for _, fieldID := range fieldIDs {
|
|
// get field schema from fieldID
|
|
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
|
|
if err != nil {
|
|
// field might have been dropped, skip it and continue processing
|
|
// other fields in the same column group
|
|
mlog.Info(ctx, "skip binlog for dropped field", mlog.FieldFieldID(fieldID), mlog.String("name", schema.GetName()))
|
|
continue
|
|
}
|
|
if _, ok := indexedFields[fieldID]; !ok {
|
|
hasIndex = false
|
|
}
|
|
|
|
// missing mapping, shall be "0" group for storage v2
|
|
if fieldSchema == nil {
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
segMemoryLoadingSize += binlogSize
|
|
}
|
|
legacyNilSchema = true
|
|
break
|
|
}
|
|
|
|
supportInterimIndexDataType = supportInterimIndexDataType || SupportInterimIndexDataType(fieldSchema.GetDataType())
|
|
isVectorType = isVectorType && typeutil.IsVectorType(fieldSchema.GetDataType())
|
|
mmapEnabled = mmapEnabled && isDataMmapEnable(fieldSchema)
|
|
containsTimestampField = containsTimestampField || DoubleMemorySystemField(fieldSchema.GetFieldID())
|
|
doubleMomoryDataField = doubleMomoryDataField || DoubleMemoryDataType(fieldSchema.GetDataType())
|
|
}
|
|
// legacy v2 segment without children
|
|
if legacyNilSchema {
|
|
continue
|
|
}
|
|
|
|
if !hasIndex {
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
interimIndexEnable := multiplyFactor.EnableInterminSegmentIndex && !isGrowingMmapEnable() && supportInterimIndexDataType
|
|
if interimIndexEnable {
|
|
segMemoryLoadingSize += uint64(float64(binlogSize) * multiplyFactor.tempSegmentIndexFactor)
|
|
}
|
|
}
|
|
}
|
|
|
|
if isVectorType {
|
|
mmapVectorField := paramtable.Get().QueryNodeCfg.MmapVectorField.GetAsBool()
|
|
if mmapVectorField {
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
segDiskLoadingSize += binlogSize
|
|
}
|
|
} else {
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
segMemoryLoadingSize += binlogSize
|
|
}
|
|
}
|
|
continue
|
|
}
|
|
|
|
// timestamp field double in InsertRecord & TimestampIndex
|
|
if containsTimestampField {
|
|
timestampSize := lo.SumBy(fieldBinlog.GetBinlogs(), func(binlog *datapb.Binlog) int64 {
|
|
return binlog.GetEntriesNum() * 4
|
|
})
|
|
segMemoryLoadingSize += 2 * uint64(timestampSize)
|
|
}
|
|
|
|
if !mmapEnabled {
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
segMemoryLoadingSize += binlogSize
|
|
if doubleMomoryDataField {
|
|
segMemoryLoadingSize += binlogSize
|
|
}
|
|
}
|
|
} else {
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
segDiskLoadingSize += uint64(getBinlogDataMemorySize(fieldBinlog))
|
|
}
|
|
}
|
|
}
|
|
|
|
// PART 2.5: external segment adjustments
|
|
//
|
|
// External segments carry pre-computed MemorySize in fake binlogs (from
|
|
// DataNode Take sampling). Adjust the memory estimate for two external-
|
|
// specific behaviors:
|
|
// 1. Non-lazy path: apply externalRawDataFactor to cover the peak
|
|
// transient memory during download + decompress + Arrow deserialize
|
|
// (normal packed segments do not have this peak because their
|
|
// binlogs are already in Arrow IPC format).
|
|
// 2. Full-lazy path (all external fields warmup=disable): no eager
|
|
// load, so subtract the raw data size that PART 2 added.
|
|
// Also propagate EstimatedBytesPerRow to the C++ ManifestGroupTranslator
|
|
// so the tiered-cache layer sizes chunks correctly.
|
|
if typeutil.IsExternalCollection(schema) && loadInfo.GetNumOfRows() > 0 {
|
|
var fakeBinlogMemSize int64
|
|
for _, fb := range loadInfo.BinlogPaths {
|
|
fakeBinlogMemSize += getBinlogDataMemorySize(fb)
|
|
}
|
|
loadInfo.EstimatedBytesPerRow = fakeBinlogMemSize / loadInfo.GetNumOfRows()
|
|
|
|
if isExternalCollectionLazyLoad(schema) {
|
|
// Full-lazy → zero eager load. Undo PART 2's rawSize addition.
|
|
// Safety factor does not apply: no peak to cover.
|
|
if segMemoryLoadingSize >= uint64(fakeBinlogMemSize) {
|
|
segMemoryLoadingSize -= uint64(fakeBinlogMemSize)
|
|
} else {
|
|
segMemoryLoadingSize = 0
|
|
}
|
|
} else if factor := multiplyFactor.externalRawDataFactor; factor > 1.0 {
|
|
// Non-lazy → add peak margin on top of rawSize that PART 2 added.
|
|
segMemoryLoadingSize += uint64(float64(fakeBinlogMemSize) * (factor - 1.0))
|
|
}
|
|
}
|
|
|
|
// PART 3: calculate size of stats data
|
|
// stats data isn't managed by the caching layer, so its size should always be included,
|
|
// regardless of the tiered eviction value
|
|
for _, fieldBinlog := range loadInfo.Statslogs {
|
|
segMemoryLoadingSize += uint64(getBinlogDataMemorySize(fieldBinlog))
|
|
}
|
|
|
|
// PART 4: calculate size of delete data
|
|
// delete data isn't managed by the caching layer, so its size should always be included,
|
|
// regardless of the tiered eviction value
|
|
for _, fieldBinlog := range loadInfo.Deltalogs {
|
|
// MemorySize of filedBinlog is the actual size in memory, but we should also consider
|
|
// the memcpy from golang to cpp side, so the expansionFactor is set to 2.
|
|
expansionFactor := float64(2)
|
|
memSize := getBinlogDataMemorySize(fieldBinlog)
|
|
|
|
// Note: If MemorySize == DiskSize, it means the segment comes from Milvus 2.3,
|
|
// MemorySize is actually compressed DiskSize of deltalog, so we'll fallback to use
|
|
// deltaExpansionFactor to compromise the compression ratio.
|
|
if memSize == getBinlogDataDiskSize(fieldBinlog) {
|
|
expansionFactor = multiplyFactor.deltaDataExpansionFactor
|
|
}
|
|
segMemoryLoadingSize += uint64(float64(memSize) * expansionFactor)
|
|
}
|
|
|
|
// PART 5: calculate size of json key stats data
|
|
jsonStatsMmapEnable := paramtable.Get().QueryNodeCfg.MmapJSONStats.GetAsBool()
|
|
for _, jsonKeyStats := range loadInfo.GetJsonKeyStatsLogs() {
|
|
if jsonStatsMmapEnable {
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
segDiskLoadingSize += uint64(float64(jsonKeyStats.GetMemorySize()) * multiplyFactor.jsonKeyStatsExpansionFactor)
|
|
}
|
|
} else {
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
segMemoryLoadingSize += uint64(float64(jsonKeyStats.GetMemorySize()) * multiplyFactor.jsonKeyStatsExpansionFactor)
|
|
}
|
|
}
|
|
}
|
|
|
|
// per struct memory size, used to keep mapping between row id and element id
|
|
var structArrayOffsetsSize uint64
|
|
// PART 6: calculate size of struct array offsets
|
|
// The memory size is 4 * row_count + 4 * total_element_count
|
|
// We cannot easily get the element count, so we estimate it by the row count * 10
|
|
rowCount := uint64(loadInfo.GetNumOfRows())
|
|
for range len(schema.GetStructArrayFields()) {
|
|
structArrayOffsetsSize += 4*rowCount + 4*rowCount*10
|
|
}
|
|
|
|
// PART 7: calculate size of text index stats data
|
|
// text index data is managed by the caching layer when tiered eviction is enabled,
|
|
// so it only needs to be included when tiered eviction is disabled.
|
|
// Text match index mmap is driven by scalar_field_enable_mmap (same as raw scalar data).
|
|
// memory_size = sum of Tantivy index file sizes (same value as C++ ByteSize() after load),
|
|
// so 1.0x is the baseline; textIndexExpansionFactor allows tuning if needed.
|
|
textIndexMmapEnable := paramtable.Get().QueryNodeCfg.MmapScalarField.GetAsBool()
|
|
for _, textStats := range loadInfo.GetTextStatsLogs() {
|
|
if textIndexMmapEnable {
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
segDiskLoadingSize += uint64(float64(textStats.GetMemorySize()) * multiplyFactor.textIndexExpansionFactor)
|
|
}
|
|
} else {
|
|
if !multiplyFactor.TieredEvictionEnabled {
|
|
segMemoryLoadingSize += uint64(float64(textStats.GetMemorySize()) * multiplyFactor.textIndexExpansionFactor)
|
|
}
|
|
}
|
|
}
|
|
|
|
return &ResourceUsage{
|
|
MemorySize: segMemoryLoadingSize + indexMemorySize + structArrayOffsetsSize,
|
|
DiskSize: segDiskLoadingSize,
|
|
MmapFieldCount: mmapFieldCount,
|
|
FieldGpuMemorySize: fieldGpuMemorySize,
|
|
}, nil
|
|
}
|
|
|
|
func DoubleMemoryDataType(dataType schemapb.DataType) bool {
|
|
return dataType == schemapb.DataType_String ||
|
|
dataType == schemapb.DataType_VarChar ||
|
|
dataType == schemapb.DataType_JSON
|
|
}
|
|
|
|
func DoubleMemorySystemField(fieldID int64) bool {
|
|
return fieldID == common.TimeStampField
|
|
}
|
|
|
|
func SupportInterimIndexDataType(dataType schemapb.DataType) bool {
|
|
return dataType == schemapb.DataType_FloatVector ||
|
|
dataType == schemapb.DataType_SparseFloatVector ||
|
|
dataType == schemapb.DataType_Float16Vector ||
|
|
dataType == schemapb.DataType_BFloat16Vector
|
|
}
|
|
|
|
func (loader *segmentLoader) getFieldType(collectionID, fieldID int64) (schemapb.DataType, error) {
|
|
collection := loader.manager.Collection.Get(collectionID)
|
|
if collection == nil {
|
|
return 0, merr.WrapErrCollectionNotFound(collectionID)
|
|
}
|
|
|
|
for _, field := range collection.Schema().GetFields() {
|
|
if field.GetFieldID() == fieldID {
|
|
return field.GetDataType(), nil
|
|
}
|
|
}
|
|
|
|
for _, structField := range collection.Schema().GetStructArrayFields() {
|
|
if structField.GetFieldID() == fieldID {
|
|
return schemapb.DataType_ArrayOfStruct, nil
|
|
}
|
|
for _, subField := range structField.GetFields() {
|
|
if subField.GetFieldID() == fieldID {
|
|
return subField.GetDataType(), nil
|
|
}
|
|
}
|
|
}
|
|
|
|
return 0, merr.WrapErrFieldNotFound(fieldID)
|
|
}
|
|
|
|
func (loader *segmentLoader) LoadIndex(ctx context.Context,
|
|
seg Segment,
|
|
loadInfo *querypb.SegmentLoadInfo,
|
|
version int64,
|
|
) error {
|
|
segment, ok := seg.(*LocalSegment)
|
|
if !ok {
|
|
return merr.WrapErrParameterInvalid("LocalSegment", fmt.Sprintf("%T", seg))
|
|
}
|
|
// Filter out LOADING segments only
|
|
// use None to avoid loaded check
|
|
infos := loader.prepare(ctx, commonpb.SegmentState_SegmentStateNone, loadInfo)
|
|
defer loader.unregister(infos...)
|
|
|
|
indexInfo := lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) *querypb.SegmentLoadInfo {
|
|
info = typeutil.Clone(info)
|
|
// remain binlog paths whose field id is in index infos to estimate resource usage correctly
|
|
indexFields := typeutil.NewSet(lo.Map(info.GetIndexInfos(), func(indexInfo *querypb.FieldIndexInfo, _ int) int64 { return indexInfo.GetFieldID() })...)
|
|
var binlogPaths []*datapb.FieldBinlog
|
|
for _, binlog := range info.GetBinlogPaths() {
|
|
if indexFields.Contain(binlog.GetFieldID()) {
|
|
binlogPaths = append(binlogPaths, binlog)
|
|
}
|
|
}
|
|
info.BinlogPaths = binlogPaths
|
|
info.Deltalogs = nil
|
|
info.Statslogs = nil
|
|
return info
|
|
})
|
|
requestResourceResult, err := loader.requestResource(ctx, indexInfo...)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer loader.freeRequestResource(requestResourceResult)
|
|
|
|
mlog.Info(context.TODO(), "segment loader start to load index", mlog.Int("segmentNumAfterFilter", len(infos)))
|
|
metrics.QueryNodeLoadSegmentConcurrency.WithLabelValues(paramtable.GetStringNodeID(), "LoadIndex").Inc()
|
|
defer metrics.QueryNodeLoadSegmentConcurrency.WithLabelValues(paramtable.GetStringNodeID(), "LoadIndex").Dec()
|
|
|
|
tr := timerecord.NewTimeRecorder("segmentLoader.LoadIndex")
|
|
defer metrics.QueryNodeLoadIndexLatency.WithLabelValues(paramtable.GetStringNodeID()).Observe(float64(tr.ElapseSpan().Milliseconds()))
|
|
for _, loadInfo := range infos {
|
|
for _, info := range loadInfo.GetIndexInfos() {
|
|
if len(info.GetIndexFilePaths()) == 0 {
|
|
mlog.Warn(context.TODO(), "failed to add index for segment, index file list is empty, the segment may be too small")
|
|
return merr.WrapErrIndexNotFound("index file list empty")
|
|
}
|
|
|
|
err := loader.loadFieldIndex(ctx, segment, info)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "failed to load index for segment", mlog.Err(err))
|
|
return err
|
|
}
|
|
}
|
|
loader.notifyLoadFinish(loadInfo)
|
|
}
|
|
|
|
return loader.waitSegmentLoadDone(ctx, commonpb.SegmentState_SegmentStateNone, []int64{loadInfo.GetSegmentID()}, version)
|
|
}
|
|
|
|
func (loader *segmentLoader) ReopenSegments(ctx context.Context,
|
|
loadInfos []*querypb.SegmentLoadInfo,
|
|
) error {
|
|
// Filter out LOADING segments only
|
|
// use None to avoid loaded check
|
|
infos := loader.prepare(ctx, commonpb.SegmentState_SegmentStateNone, loadInfos...)
|
|
defer loader.unregister(infos...)
|
|
|
|
// use full resource in case of whole segment reopen
|
|
// TODO use calculated resource from segcore after supported
|
|
requestResourceResult, err := loader.requestResource(ctx, infos...)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "reopen segment request resource failed", mlog.Err(err))
|
|
return err
|
|
}
|
|
defer loader.freeRequestResource(requestResourceResult)
|
|
|
|
for _, info := range infos {
|
|
segment := loader.manager.Segment.GetSealed(info.GetSegmentID())
|
|
if segment == nil {
|
|
mlog.Warn(context.TODO(), "failed to reopen segment, segment not loaded", mlog.Int64("segmentID", info.GetSegmentID()))
|
|
continue
|
|
}
|
|
collection := loader.manager.Collection.Get(info.GetCollectionID())
|
|
if collection != nil {
|
|
configureUseTakeForOutput(info, collection.Schema())
|
|
}
|
|
|
|
err := segment.Reopen(ctx, info)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "failed to reopen segment", mlog.Int64("segmentID", info.GetSegmentID()), mlog.Err(err))
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func getBinlogDataDiskSize(fieldBinlog *datapb.FieldBinlog) int64 {
|
|
fieldSize := int64(0)
|
|
for _, binlog := range fieldBinlog.Binlogs {
|
|
fieldSize += binlog.GetLogSize()
|
|
}
|
|
|
|
return fieldSize
|
|
}
|
|
|
|
func getBinlogDataMemorySize(fieldBinlog *datapb.FieldBinlog) int64 {
|
|
fieldSize := int64(0)
|
|
for _, binlog := range fieldBinlog.Binlogs {
|
|
fieldSize += binlog.GetMemorySize()
|
|
}
|
|
|
|
return fieldSize
|
|
}
|
|
|
|
func gpuIndexRequiresGpu(indexParams []*commonpb.KeyValuePair) bool {
|
|
indexParamMap := funcutil.KeyValuePair2Map(indexParams)
|
|
indexType := indexParamMap[common.IndexTypeKey]
|
|
|
|
switch indexType {
|
|
case "GPU_CAGRA", "GPU_CUVS_CAGRA":
|
|
case "GPU_BRUTE_FORCE", "GPU_CUVS_BRUTE_FORCE",
|
|
"GPU_IVF_FLAT", "GPU_CUVS_IVF_FLAT",
|
|
"GPU_IVF_PQ", "GPU_CUVS_IVF_PQ":
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
|
|
err := indexparams.AppendPrepareLoadParams(paramtable.Get(), indexParamMap)
|
|
if err != nil {
|
|
mlog.Warn(context.TODO(), "failed to append prepare load params for gpu index resource check",
|
|
mlog.String("indexType", indexType),
|
|
mlog.Err(err))
|
|
}
|
|
|
|
adaptForCPU, ok := indexParamMap["adapt_for_cpu"]
|
|
if ok {
|
|
enabled, err := strconv.ParseBool(adaptForCPU)
|
|
if err == nil && enabled {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func checkSegmentGpuMemSize(fieldGpuMemSizeList []uint64, OverloadedMemoryThresholdPercentage float32) error {
|
|
gpuInfos, err := hardware.GetAllGPUMemoryInfo()
|
|
if err != nil {
|
|
if len(fieldGpuMemSizeList) == 0 {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
var usedGpuMem []uint64
|
|
var maxGpuMemSize []uint64
|
|
for _, gpuInfo := range gpuInfos {
|
|
usedGpuMem = append(usedGpuMem, gpuInfo.TotalMemory-gpuInfo.FreeMemory)
|
|
maxGpuMemSize = append(maxGpuMemSize, uint64(float32(gpuInfo.TotalMemory)*OverloadedMemoryThresholdPercentage))
|
|
}
|
|
currentGpuMem := usedGpuMem
|
|
for _, fieldGpuMem := range fieldGpuMemSizeList {
|
|
var minId int = -1
|
|
var minGpuMem uint64 = math.MaxUint64
|
|
for i := int(0); i < len(gpuInfos); i++ {
|
|
GpuiMem := currentGpuMem[i] + fieldGpuMem
|
|
if GpuiMem < maxGpuMemSize[i] && GpuiMem < minGpuMem {
|
|
minId = i
|
|
minGpuMem = GpuiMem
|
|
}
|
|
}
|
|
if minId == -1 {
|
|
mlog.Warn(context.TODO(), "load segment failed, GPU OOM if loaded",
|
|
mlog.String("resourceType", "GPU"),
|
|
mlog.Uint64("gpuMemUsageBytes", fieldGpuMem),
|
|
mlog.Any("usedGpuMemBytes", usedGpuMem),
|
|
mlog.Any("maxGpuMemBytes", maxGpuMemSize),
|
|
)
|
|
return merr.WrapErrSegmentRequestResourceFailed("GPU")
|
|
}
|
|
currentGpuMem[minId] = minGpuMem
|
|
}
|
|
return nil
|
|
}
|