Files
milvus/internal/querynodev2/segments/segment_loader.go
T
sijie-ni-0214 24db48a7e6 enhance: add QueryCoord autoscale load resource precheck (#51218)
## Summary

issue: #51217

Load requests introduce resource demand before QueryNode waterlines
reflect the new data. This PR adds optional QueryCoord resource
prechecks before LoadCollection/LoadPartitions configuration broadcasts
and publishes cumulative demand counters for external autoscaling
decisions.

Design: [Autoscale-Aware Load Resource Admission in
QueryCoord](https://github.com/sijie-ni-0214/milvus/blob/querycoord-autoscale-kernel/docs/design-docs/design_docs/20260721-querycoord-autoscale-load-resource-precheck.md).

## Admission behavior

- Estimate incremental final sealed-segment memory and local-disk demand
per resource group from the requested and current load configurations.
- With tiered eviction disabled, compare demand with memory availability
and each Worker's configured `LOCAL_STORAGE_SIZE` quota minus cached
local-directory usage after admission thresholds. Optionally admit
shortages within configured global autoscale upper bounds.
- With tiered eviction enabled, retain demand estimation and skip QN
capacity sampling, node availability checks, and global autoscale-limit
checks. Existing evictable caches make physical usage unsuitable for
this admission decision; Worker QN retains its loading-stage guard.
- Fail open on metadata, estimation, and capacity-metrics failures.
Capacity-metrics failures preserve successfully estimated demand;
metadata or estimation failures leave no demand to publish.
- Return a resource-insufficient status with `suggested_expand_percent`
when a complete capacity evaluation rejects the request. Proxy preserves
the status and ExtraInfo.
- Admission remains best-effort, without reservations or accounting for
concurrent accepted requests. QueryCoord does not subtract Worker
in-flight load reservations; the Worker guard remains authoritative. The
existing replica node-count prerequisite remains; node provisioning and
scale-from-zero are outside this PR.

## Load demand counters

Expose cumulative counters labeled by resource group:

- `milvus_querycoord_load_demand_memory_bytes`
- `milvus_querycoord_load_demand_disk_bytes`

Counters increase only when demand estimation succeeds, the request
proceeds, and its configuration broadcast succeeds. This includes tiered
loads and loads admitted through capacity-metrics fail-open. Disabled
prechecks, unchanged configurations, zero incremental demand, estimation
failures, explicit rejections, and failed broadcasts do not increase
counters. Load completion and release do not decrement them.

Series start at zero for existing resource groups and are removed on RG
deletion. The dashboard displays cumulative values in MiB; the control
plane computes interval increments, handles counter resets, and combines
demand with QN waterlines and provisioning state. These counters are
observations, not durable reservations. Internal handoff and balancing
tasks do not contribute to them.

## Resource estimation and sampling

- Share segment/index estimation in
`internal/util/segcore/loadresource`, retaining separate final-resource
and loading-peak APIs. Worker QN uses its CGo pool; QueryCoord estimates
sequentially within each request.
- Fetch recovery metadata, full segment metadata, and batched index
metadata. Support Storage V3 resource summaries and dropped-field
metadata. Filter indexes by loaded fields, retaining multiple JSON-path
indexes on one field.
- Apply the load-request collection mmap inheritance logic to cloned
expected/current schemas while preserving field overrides. The final
estimator applies cache ratios only with tiered eviction enabled.
Delegator-side Bloom/BM25 resources are outside the Worker estimate.
- Deduplicate QN resource sampling by node ID, use at most 16 concurrent
requests and one attempt per node, and bound the entire sampling round
to ten seconds. Reuse the snapshot for RG and global calculations.
Exclude embedded SQNs from global capacity and sampling; empty
eligible-node sets are unevaluable and fail open.
- Worker QNs report `LOCAL_STORAGE_SIZE` and cached local-directory
usage through `system_info`. QueryCoord uses these byte values for disk
availability, RG suggestions, and the global disk-capacity sum;
filesystem-wide disk metrics remain for monitoring. Shared filesystems
require per-QN quotas consistent with the shared capacity.

## Configuration

- `queryCoord.autoscale.precheckEnabled`: enable demand estimation and
precheck; default `false`.
- `queryCoord.autoscale.enabled`: allow admission within global
autoscale bounds when current capacity is insufficient; default `false`,
effective only with precheck enabled and tiered eviction disabled.
- `queryCoord.autoscale.maxMemoryLimit`: maximum global Worker QN
physical memory capacity after scaling, in GiB; default `0` means no
additional memory autoscale capacity.
- `queryCoord.autoscale.maxDiskLimit`: maximum sum of Worker QN
`LOCAL_STORAGE_SIZE` quotas after scaling, in GiB; default `0` means no
additional disk autoscale capacity. Both limits use `1 GiB = 1024^3
bytes`.

For non-tiered deployments, configure QueryCoord's memory high-water
level no higher than the Worker admission threshold. This relationship
is a rollout prerequisite and is not enforced by this PR.

## Validation

Local autoscale estimator and targeted QueryCoord tests passed,
including multiple-index filtering, collection mmap overrides, the
ten-second sampling timeout, empty-node fail-open, embedded-SQN
exclusion, tiered demand preservation, capacity-metrics failure
behavior, and local-storage quota accounting for QueryNodes sharing a
filesystem.

Both load callbacks are tested for precheck rejection, broadcast
failure, successful counter updates, tiered loads, and capacity-metrics
failures followed by successful or failed broadcasts. Tests assert
counters remain unchanged during Broadcast and increment only after it
succeeds. These are local unit/callback checks, not end-to-end external
autoscaler validation.

Local `make static-check` passed across core, pkg, client, and go_client
with zero issues, and segcore error-boundary checks passed. Targeted
QueryCoord and QueryNode tests, YAML consistency, and `make build-go`
passed. External autoscaler behavior remains subject to CI and rollout
validation.

---------

Signed-off-by: sijie-ni-0214 <[email protected]>
2026-09-29 01:47:36 +08:00

2130 lines
76 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/metastore/kv/binlog"
"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/segcore/loadresource"
"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/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
// GetLocalDiskUsage returns the cached size of the local storage directory.
GetLocalDiskUsage() (int64, error)
// ReopenSegments update segment data according to new load info.
ReopenSegments(ctx context.Context,
loadInfos []*querypb.SegmentLoadInfo,
) error
}
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 when tiered eviction is disabled. With tiered eviction enabled,
// the caching layer reserves transient loading overhead from the sampled
// external row size, so applying this factor in Go would reserve the same
// raw-data memory twice. 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) GetLocalDiskUsage() (int64, error) {
return loader.duf.GetDiskUsage()
}
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
if err := prepareIndexLoadParams(loadInfo.GetIndexInfos()); err != nil {
return nil, err
}
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)
if localSegment, ok := segment.(*LocalSegment); ok {
localSegment.compactLoadInfoForRuntime()
}
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.
// Match the writer's primary storage root; local stats do not use MinIO's prefix.
rootPath := binlog.GetRootPath()
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
}
}
relatedDataSize := calculateSegmentLogSize(segment.LoadInfo())
segment.relatedDataSize.Store(relatedDataSize)
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) 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) 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(ctx, pkField.DataType, paths, opts...)
if err != nil {
return err
}
return readDeltaRecords(reader)
}
// Manifest-backed delta loading is shared by the parent segment and by
// compact-to child manifests carried as a load-time delete overlay.
readManifestDeltas := func(manifestPath string) error {
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(
ctx,
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(
ctx,
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
}
}
return nil
}
if manifestPath := loadInfo.GetManifestPath(); manifestPath != "" && !useExplicitDeltalogs {
if err := readManifestDeltas(manifestPath); err != nil {
return err
}
} else {
// V1: delta data referenced by Deltalogs entries
paths := make([]string, 0)
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
}
}
// Child manifests are loaded after the parent delete source so all delete
// records are folded into the same DeltaData before segcore sees the segment.
for _, manifestPath := range loadInfo.GetChildManifestPaths() {
if err := readManifestDeltas(manifestPath); 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(),
// External collections may reference an s3:// source even when the
// primary storage is local, so the connection cap still applies.
MaxConnections: uint32(params.MinioCfg.MaxConnections.GetAsInt()),
}
}
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(),
MaxConnections: uint32(params.MinioCfg.MaxConnections.GetAsInt()),
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(),
jsonKeyStatsExpansionFactor: paramtable.Get().QueryNodeCfg.JSONKeyStatsExpansionFactor.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) {
options := loadresource.DefaultSegmentFinalEstimateOptions()
options.DeltaDataExpansionFactor = multiplyFactor.deltaDataExpansionFactor
options.JSONKeyStatsExpansionFactor = multiplyFactor.jsonKeyStatsExpansionFactor
options.TextIndexExpansionFactor = multiplyFactor.textIndexExpansionFactor
options.TieredEvictionEnabled = multiplyFactor.TieredEvictionEnabled
options.TieredEvictableMemoryCacheRatio = multiplyFactor.TieredEvictableMemoryCacheRatio
options.TieredEvictableDiskCacheRatio = multiplyFactor.TieredEvictableDiskCacheRatio
estimate, err := loadresource.EstimateSegmentFinalResource(context.Background(), schema, loadInfo, options, func(fn func() error) error {
_, err := GetDynamicPool().Submit(func() (any, error) {
return nil, fn()
}).Await()
return err
})
if err != nil {
return nil, err
}
mlog.Debug(context.TODO(), "estimate logical resource usage result",
mlog.Int64("segmentID", loadInfo.GetSegmentID()),
mlog.Uint64("memorySize", estimate.MemoryBytes),
mlog.Uint64("diskSize", estimate.DiskBytes),
)
return &ResourceUsage{
MemorySize: estimate.MemoryBytes,
DiskSize: estimate.DiskBytes,
}, 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) {
options := loadresource.DefaultSegmentLoadingEstimateOptions()
options.DeltaDataExpansionFactor = multiplyFactor.deltaDataExpansionFactor
options.JSONKeyStatsExpansionFactor = multiplyFactor.jsonKeyStatsExpansionFactor
options.TextIndexExpansionFactor = multiplyFactor.textIndexExpansionFactor
options.TieredEvictionEnabled = multiplyFactor.TieredEvictionEnabled
options.EnableInterimSegmentIndex = multiplyFactor.EnableInterminSegmentIndex
options.TempSegmentIndexFactor = multiplyFactor.tempSegmentIndexFactor
options.ExternalRawDataFactor = multiplyFactor.externalRawDataFactor
estimate, err := loadresource.EstimateSegmentLoadingResource(context.Background(), schema, loadInfo, options, func(fn func() error) error {
_, err := GetDynamicPool().Submit(func() (any, error) {
return nil, fn()
}).Await()
return err
})
if err != nil {
return nil, err
}
return &ResourceUsage{
MemorySize: estimate.MemoryBytes,
DiskSize: estimate.DiskBytes,
MmapFieldCount: estimate.MmapFieldCount,
FieldGpuMemorySize: estimate.FieldGPUMemoryBytes,
}, nil
}
// prepareIndexLoadParams injects QueryNode-local index load parameters into each
// index's IndexParams in place. These params (e.g. DISKANN num_load_thread) are
// derived from local QueryNode resources/config and are never persisted in the
// index metadata, so they must be re-injected on every load path before the load
// info reaches segcore. Both full-load (Load) and Reopen call this; skipping it
// on Reopen was the root cause of issue #51249 (segcore asserts
// "param num_load_thread is empty" while loading a DISKANN index).
func prepareIndexLoadParams(indexInfos []*querypb.FieldIndexInfo) error {
for _, indexInfo := range indexInfos {
if indexInfo == nil {
continue
}
indexParams := funcutil.KeyValuePair2Map(indexInfo.GetIndexParams())
// 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 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 err
}
indexInfo.IndexParams = funcutil.Map2KeyValuePair(indexParams)
}
return nil
}
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
}