Files
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
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

1849 lines
58 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 "futures/future_c.h"
#include "segcore/collection_c.h"
#include "segcore/plan_c.h"
#include "segcore/segment_c.h"
#include "common/init_c.h"
*/
import "C"
import (
"context"
"fmt"
"strings"
"sync"
"time"
"unsafe"
"github.com/cockroachdb/errors"
"go.opentelemetry.io/otel"
"go.uber.org/atomic"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/querynodev2/pkoracle"
"github.com/milvus-io/milvus/internal/querynodev2/segments/state"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/internal/util/indexparamcheck"
"github.com/milvus-io/milvus/internal/util/segcore"
"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/cgopb"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/proto/segcorepb"
"github.com/milvus-io/milvus/pkg/v3/util/contextutil"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/indexparams"
"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/timerecord"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
type SegmentType = commonpb.SegmentState
const (
SegmentTypeGrowing = commonpb.SegmentState_Growing
SegmentTypeSealed = commonpb.SegmentState_Sealed
)
var ErrSegmentUnhealthy = errors.New("segment unhealthy")
// IndexedFieldInfo contains binlog info of vector field
type IndexedFieldInfo struct {
FieldBinlog *datapb.FieldBinlog
IndexInfo *querypb.FieldIndexInfo
IsLoaded bool
}
type baseSegment struct {
collection *Collection
version *atomic.Int64
segmentType SegmentType
pkCandidate pkoracle.Candidate // PK candidate: BloomFilterSet for regular collections, ExternalSegmentCandidate for external collections
loadInfo *atomic.Pointer[querypb.SegmentLoadInfo]
skipGrowingBF bool // Skip generating or maintaining BF for growing segments; deletion checks will be handled in segcore.
channel metautil.Channel
resourceUsageCache *atomic.Pointer[ResourceUsage]
needUpdatedVersion *atomic.Int64 // only for lazy load mode update index
}
func newBaseSegment(collection *Collection, segmentType SegmentType, version int64, loadInfo *querypb.SegmentLoadInfo) (baseSegment, error) {
channel, err := metautil.ParseChannel(loadInfo.GetInsertChannel(), channelMapper)
if err != nil {
return baseSegment{}, err
}
bs := baseSegment{
collection: collection,
loadInfo: atomic.NewPointer[querypb.SegmentLoadInfo](loadInfo),
version: atomic.NewInt64(version),
segmentType: segmentType,
pkCandidate: pkoracle.NewBloomFilterSet(loadInfo.GetSegmentID(), loadInfo.GetPartitionID(), segmentType),
channel: channel,
skipGrowingBF: segmentType == SegmentTypeGrowing && paramtable.Get().QueryNodeCfg.SkipGrowingSegmentBF.GetAsBool(),
resourceUsageCache: atomic.NewPointer[ResourceUsage](nil),
needUpdatedVersion: atomic.NewInt64(0),
}
return bs, nil
}
// ID returns the identity number.
func (s *baseSegment) ID() int64 {
return s.loadInfo.Load().GetSegmentID()
}
func (s *baseSegment) Collection() int64 {
return s.loadInfo.Load().GetCollectionID()
}
func (s *baseSegment) GetCollection() *Collection {
return s.collection
}
func (s *baseSegment) Partition() int64 {
return s.loadInfo.Load().GetPartitionID()
}
func (s *baseSegment) DatabaseName() string {
return s.collection.GetDBName()
}
func (s *baseSegment) ResourceGroup() string {
return s.collection.GetResourceGroup()
}
func (s *baseSegment) Shard() metautil.Channel {
return s.channel
}
func (s *baseSegment) Type() SegmentType {
return s.segmentType
}
func (s *baseSegment) Level() datapb.SegmentLevel {
return s.loadInfo.Load().GetLevel()
}
func (s *baseSegment) IsSorted() bool {
return s.loadInfo.Load().GetIsSorted()
}
func (s *baseSegment) StartPosition() *msgpb.MsgPosition {
return s.loadInfo.Load().GetStartPosition()
}
func (s *baseSegment) Version() int64 {
return s.version.Load()
}
func (s *baseSegment) CASVersion(old, newVersion int64) bool {
return s.version.CompareAndSwap(old, newVersion)
}
func (s *baseSegment) LoadInfo() *querypb.SegmentLoadInfo {
return s.loadInfo.Load()
}
func (s *baseSegment) SetPKCandidate(candidate pkoracle.Candidate) {
s.pkCandidate = candidate
}
// PkCandidateExist implements pkoracle.Candidate — reports whether PK data has been loaded.
func (s *baseSegment) PkCandidateExist() bool {
return s.pkCandidate != nil && s.pkCandidate.PkCandidateExist()
}
// UpdatePkCandidate feeds new primary keys into the PK candidate.
func (s *baseSegment) UpdatePkCandidate(pks []storage.PrimaryKey) {
if s.skipGrowingBF {
return
}
if s.pkCandidate != nil {
s.pkCandidate.UpdatePkCandidate(pks)
}
}
// Stats implements pkoracle.Candidate — returns PK statistics (min/max PK).
func (s *baseSegment) Stats() *storage.PkStatistics {
if s.pkCandidate != nil {
return s.pkCandidate.Stats()
}
return nil
}
// Charge implements pkoracle.Candidate — charges memory resources.
func (s *baseSegment) Charge() {
if s.pkCandidate != nil {
s.pkCandidate.Charge()
}
}
// Refund implements pkoracle.Candidate — releases memory resources.
func (s *baseSegment) Refund() {
if s.pkCandidate != nil {
s.pkCandidate.Refund()
}
}
type bm25StatsHolder struct {
mu sync.RWMutex
stats map[int64]*storage.BM25Stats
}
func newBM25StatsHolder() *bm25StatsHolder {
return &bm25StatsHolder{
stats: make(map[int64]*storage.BM25Stats),
}
}
func (h *bm25StatsHolder) UpdateBM25Stats(stats map[int64]*storage.BM25Stats) {
if h == nil {
return
}
h.mu.Lock()
defer h.mu.Unlock()
for fieldID, new := range stats {
if new == nil {
continue
}
if current, ok := h.stats[fieldID]; ok {
current.Merge(new)
} else {
h.stats[fieldID] = new.Clone()
}
}
}
func (h *bm25StatsHolder) GetBM25Stats() map[int64]*storage.BM25Stats {
if h == nil {
return map[int64]*storage.BM25Stats{}
}
h.mu.RLock()
defer h.mu.RUnlock()
stats := cloneBM25StatsMap(h.stats)
if stats == nil {
return map[int64]*storage.BM25Stats{}
}
return stats
}
func cloneBM25StatsMap(stats map[int64]*storage.BM25Stats) map[int64]*storage.BM25Stats {
if len(stats) == 0 {
return nil
}
cloned := make(map[int64]*storage.BM25Stats, len(stats))
for fieldID, stat := range stats {
if stat != nil {
cloned[fieldID] = stat.Clone()
}
}
return cloned
}
// MayPkExist returns true if the given PK exists in the PK range and being positive through the bloom filter,
// false otherwise,
// may returns true even the PK doesn't exist actually
func (s *baseSegment) MayPkExist(pk *storage.LocationsCache) bool {
if s.skipGrowingBF {
return true
}
if s.pkCandidate == nil {
return true // No candidate, assume PK might exist
}
return s.pkCandidate.MayPkExist(pk)
}
func (s *baseSegment) GetMinPk() *storage.PrimaryKey {
if s.pkCandidate != nil {
if stats := s.pkCandidate.Stats(); stats != nil {
return &stats.MinPK
}
}
return nil
}
func (s *baseSegment) GetMaxPk() *storage.PrimaryKey {
if s.pkCandidate != nil {
if stats := s.pkCandidate.Stats(); stats != nil {
return &stats.MaxPK
}
}
return nil
}
func (s *baseSegment) BatchPkExist(lc *storage.BatchLocationsCache) []bool {
if s.skipGrowingBF {
allPositive := make([]bool, lc.Size())
for i := 0; i < lc.Size(); i++ {
allPositive[i] = true
}
return allPositive
}
if s.pkCandidate == nil {
allPositive := make([]bool, lc.Size())
for i := 0; i < lc.Size(); i++ {
allPositive[i] = true
}
return allPositive
}
return s.pkCandidate.BatchPkExist(lc)
}
// ResourceUsageEstimate returns the final estimated resource usage of the segment.
func (s *baseSegment) ResourceUsageEstimate() ResourceUsage {
if s.segmentType == SegmentTypeGrowing {
// Growing segment cannot do resource usage estimate.
return ResourceUsage{}
}
cache := s.resourceUsageCache.Load()
if cache != nil {
return *cache
}
usage, err := estimateLogicalResourceUsageOfSegment(s.collection.Schema(), s.LoadInfo(), resourceEstimateFactor{
deltaDataExpansionFactor: paramtable.Get().QueryNodeCfg.DeltaDataExpansionRate.GetAsFloat(),
TieredEvictionEnabled: paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool(),
TieredEvictableMemoryCacheRatio: paramtable.Get().QueryNodeCfg.TieredEvictableMemoryCacheRatio.GetAsFloat(),
TieredEvictableDiskCacheRatio: paramtable.Get().QueryNodeCfg.TieredEvictableDiskCacheRatio.GetAsFloat(),
})
if err != nil {
// Should never failure, if failed, segment should never be loaded.
mlog.Warn(context.TODO(), "unreachable: failed to get resource usage estimate of segment", mlog.Err(err), mlog.FieldCollectionID(s.Collection()), mlog.FieldSegmentID(s.ID()))
return ResourceUsage{}
}
s.resourceUsageCache.Store(usage)
return *usage
}
func (s *baseSegment) NeedUpdatedVersion() int64 {
return s.needUpdatedVersion.Load()
}
func (s *baseSegment) SetLoadInfo(loadInfo *querypb.SegmentLoadInfo) {
s.loadInfo.Store(loadInfo)
}
func (s *baseSegment) SetNeedUpdatedVersion(version int64) {
s.needUpdatedVersion.Store(version)
}
type FieldInfo struct {
*datapb.FieldBinlog
RowCount int64
}
var _ Segment = (*LocalSegment)(nil)
// Segment is a wrapper of the underlying C-structure segment.
type LocalSegment struct {
baseSegment
*bm25StatsHolder
manager SegmentManager
ptrLock *state.LoadStateLock
ptr C.CSegmentInterface // TODO: Remove in future, after move load index into segcore package.
// always keep same with csegment.RawPtr(), for eaiser to access,
csegment segcore.CSegment
// cached results, to avoid too many CGO calls
memSize *atomic.Int64
binlogSize *atomic.Int64
rowNum *atomic.Int64
insertCount *atomic.Int64
deltaMut sync.Mutex
lastDeltaTimestamp *atomic.Uint64
fields *typeutil.ConcurrentMap[int64, *FieldInfo]
fieldIndexes *typeutil.ConcurrentMap[int64, *IndexedFieldInfo] // indexID -> IndexedFieldInfo
fieldJSONStats map[int64]*querypb.JsonStatsInfo
fieldJSONStatsMu sync.RWMutex
}
func NewSegment(ctx context.Context,
collection *Collection,
manager SegmentManager,
segmentType SegmentType,
version int64,
loadInfo *querypb.SegmentLoadInfo,
) (Segment, error) {
/*
CStatus
NewSegment(CCollection collection, uint64_t segment_id, SegmentType seg_type, CSegmentInterface* newSegment);
*/
if loadInfo.GetLevel() == datapb.SegmentLevel_L0 {
return NewL0Segment(collection, segmentType, version, loadInfo)
}
base, err := newBaseSegment(collection, segmentType, version, loadInfo)
if err != nil {
return nil, err
}
var locker *state.LoadStateLock
switch segmentType {
case SegmentTypeSealed:
locker = state.NewLoadStateLock(state.LoadStateOnlyMeta)
case SegmentTypeGrowing:
locker = state.NewLoadStateLock(state.LoadStateDataLoaded)
default:
return nil, merr.WrapErrServiceInternalMsg("illegal segment type %d when create segment %d", segmentType, loadInfo.GetSegmentID())
}
logger := mlog.With(
mlog.FieldCollectionID(loadInfo.GetCollectionID()),
mlog.FieldPartitionID(loadInfo.GetPartitionID()),
mlog.FieldSegmentID(loadInfo.GetSegmentID()),
mlog.String("segmentType", segmentType.String()),
mlog.String("level", loadInfo.GetLevel().String()),
)
var csegment segcore.CSegment
if _, err := GetDynamicPool().Submit(func() (any, error) {
var err error
csegment, err = collection.CreateCSegment(&segcore.CreateCSegmentRequest{
SegmentID: loadInfo.GetSegmentID(),
SegmentType: segmentType,
IsSorted: loadInfo.GetIsSorted(),
LoadInfo: loadInfo,
})
return nil, err
}).Await(); err != nil {
logger.Warn(ctx, "create segment failed", mlog.Err(err))
return nil, err
}
logger.Info(ctx, "create segment done")
segment := &LocalSegment{
baseSegment: base,
bm25StatsHolder: newBM25StatsHolder(),
manager: manager,
ptrLock: locker,
ptr: C.CSegmentInterface(csegment.RawPointer()),
csegment: csegment,
lastDeltaTimestamp: atomic.NewUint64(0),
fields: typeutil.NewConcurrentMap[int64, *FieldInfo](),
fieldIndexes: typeutil.NewConcurrentMap[int64, *IndexedFieldInfo](),
fieldJSONStats: make(map[int64]*querypb.JsonStatsInfo),
memSize: atomic.NewInt64(-1),
binlogSize: atomic.NewInt64(0),
rowNum: atomic.NewInt64(-1),
insertCount: atomic.NewInt64(0),
}
if err := segment.initializeSegment(); err != nil {
csegment.Release()
return nil, err
}
return segment, nil
}
func (s *LocalSegment) initializeSegment() error {
loadInfo := s.loadInfo.Load()
indexedFieldInfos, fieldBinlogs := separateIndexAndBinlog(loadInfo)
schemaHelper, _ := typeutil.CreateSchemaHelper(s.collection.Schema())
for _, info := range indexedFieldInfos {
fieldID := info.IndexInfo.FieldID
field, err := schemaHelper.GetFieldFromID(fieldID)
if err != nil {
return err
}
indexInfo := info.IndexInfo
s.fieldIndexes.Insert(indexInfo.GetIndexID(), &IndexedFieldInfo{
FieldBinlog: &datapb.FieldBinlog{
FieldID: indexInfo.GetFieldID(),
},
IndexInfo: indexInfo,
IsLoaded: false,
})
if !typeutil.IsVectorType(field.GetDataType()) && !s.HasRawData(fieldID) {
s.fields.Insert(fieldID, &FieldInfo{
FieldBinlog: info.FieldBinlog,
RowCount: loadInfo.GetNumOfRows(),
})
}
}
for _, binlogs := range fieldBinlogs {
s.fields.Insert(binlogs.FieldID, &FieldInfo{
FieldBinlog: binlogs,
RowCount: loadInfo.GetNumOfRows(),
})
}
// Update the insert count when initialize the segment and update the metrics.
s.insertCount.Store(loadInfo.GetNumOfRows())
return nil
}
// PinIfNotReleased acquires the `ptrLock` and returns true if the pointer is valid
// Provide ONLY the read lock operations,
// don't make `ptrLock` public to avoid abusing of the mutex.
func (s *LocalSegment) PinIfNotReleased() error {
if !s.ptrLock.PinIf(state.IsNotReleased) {
return merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
return nil
}
func (s *LocalSegment) Unpin() {
s.ptrLock.Unpin()
}
func (s *LocalSegment) InsertCount() int64 {
return s.insertCount.Load()
}
func (s *LocalSegment) RowNum() int64 {
// if segment is not loaded, return 0 (maybe not loaded or release by lru)
if !s.ptrLock.PinIf(state.IsDataLoaded) {
return 0
}
defer s.ptrLock.Unpin()
rowNum := s.rowNum.Load()
if rowNum < 0 {
GetDynamicPool().Submit(func() (any, error) {
rowNum = s.csegment.RowNum()
s.rowNum.Store(rowNum)
return nil, nil
}).Await()
}
return rowNum
}
func (s *LocalSegment) MemSize() int64 {
if !s.ptrLock.PinIf(state.IsNotReleased) {
return 0
}
defer s.ptrLock.Unpin()
memSize := s.memSize.Load()
if memSize < 0 {
GetDynamicPool().Submit(func() (any, error) {
memSize = s.csegment.MemSize()
s.memSize.Store(memSize)
return nil, nil
}).Await()
}
return memSize
}
func (s *LocalSegment) LastDeltaTimestamp() uint64 {
return s.lastDeltaTimestamp.Load()
}
// advanceLastDeltaTimestamp moves lastDeltaTimestamp forward to max(current, max(tss)).
// Consumers (file-level skip in segment_loader.LoadDeltaLogs, dist_handler reporting to
// QueryCoord) treat this field as a high-water-mark. Using tss[last] on unsorted batches
// underestimates the watermark, so we scan for the true max.
func (s *LocalSegment) advanceLastDeltaTimestamp(tss []typeutil.Timestamp) {
if len(tss) == 0 {
return
}
maxTs := tss[0]
for _, t := range tss[1:] {
if t > maxTs {
maxTs = t
}
}
for {
cur := s.lastDeltaTimestamp.Load()
if maxTs <= cur {
return
}
if s.lastDeltaTimestamp.CompareAndSwap(cur, maxTs) {
return
}
}
}
// UpdatePkCandidate updates the PK candidate with provided pks and charges resource.
// Overrides baseSegment.UpdatePkCandidate to handle resource charging for growing segments.
func (s *LocalSegment) UpdatePkCandidate(pks []storage.PrimaryKey) {
if s.skipGrowingBF {
return
}
s.pkCandidate.UpdatePkCandidate(pks)
// Charge resource (safe to call multiple times - only charges once)
s.pkCandidate.Charge()
}
func (s *LocalSegment) GetIndexByID(indexID int64) *IndexedFieldInfo {
info, _ := s.fieldIndexes.Get(indexID)
return info
}
func (s *LocalSegment) GetIndex(fieldID int64) []*IndexedFieldInfo {
var info []*IndexedFieldInfo
s.fieldIndexes.Range(func(key int64, value *IndexedFieldInfo) bool {
if value.IndexInfo.FieldID == fieldID {
info = append(info, value)
}
return true
})
return info
}
func (s *LocalSegment) ExistIndex(fieldID int64) bool {
contain := false
s.fieldIndexes.Range(func(key int64, value *IndexedFieldInfo) bool {
if value.IndexInfo.FieldID == fieldID {
contain = true
}
return !contain
})
return contain
}
func (s *LocalSegment) HasRawData(fieldID int64) bool {
if !s.ptrLock.PinIf(state.IsNotReleased) {
return false
}
defer s.ptrLock.Unpin()
return s.csegment.HasRawData(fieldID)
}
func (s *LocalSegment) HasFieldData(fieldID int64) bool {
if !s.ptrLock.PinIf(state.IsNotReleased) {
return false
}
defer s.ptrLock.Unpin()
return s.csegment.HasFieldData(fieldID)
}
func (s *LocalSegment) DropIndex(ctx context.Context, indexID int64) error {
if !s.ptrLock.PinIf(state.IsNotReleased) {
return merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
defer s.ptrLock.Unpin()
if indexInfo, ok := s.fieldIndexes.Get(indexID); ok {
field := typeutil.GetField(s.collection.Schema(), indexInfo.IndexInfo.FieldID)
if typeutil.IsJSONType(field.GetDataType()) {
nestedPath, err := funcutil.GetAttrByKeyFromRepeatedKV(common.JSONPathKey, indexInfo.IndexInfo.GetIndexParams())
if err != nil {
return err
}
err = s.csegment.DropJSONIndex(ctx, indexInfo.IndexInfo.FieldID, nestedPath)
if err != nil {
return err
}
} else {
err := s.csegment.DropIndex(ctx, indexInfo.IndexInfo.FieldID)
if err != nil {
return err
}
}
s.fieldIndexes.Remove(indexID)
}
return nil
}
func (s *LocalSegment) Indexes() []*IndexedFieldInfo {
var result []*IndexedFieldInfo
s.fieldIndexes.Range(func(key int64, value *IndexedFieldInfo) bool {
result = append(result, value)
return true
})
return result
}
func (s *LocalSegment) IsLazyLoad() bool {
for _, indexInfo := range s.Indexes() {
if !indexInfo.IsLoaded {
return true
}
}
return false
}
func (s *LocalSegment) ResetIndexesLazyLoad(lazyState bool) {
for _, indexInfo := range s.Indexes() {
indexInfo.IsLoaded = lazyState
}
}
func (s *LocalSegment) syncFieldIndexes(indexInfos []*querypb.FieldIndexInfo) {
isLoadedByIndexID := make(map[int64]bool)
loadedForNewIndex := !s.IsLazyLoad()
for _, index := range s.Indexes() {
isLoadedByIndexID[index.IndexInfo.GetIndexID()] = index.IsLoaded
}
indexIDs := typeutil.NewSet[int64]()
for _, indexInfo := range indexInfos {
indexID := indexInfo.GetIndexID()
indexIDs.Insert(indexID)
isLoaded, ok := isLoadedByIndexID[indexID]
if !ok {
isLoaded = loadedForNewIndex
}
s.fieldIndexes.Insert(indexID, &IndexedFieldInfo{
FieldBinlog: &datapb.FieldBinlog{
FieldID: indexInfo.GetFieldID(),
},
IndexInfo: indexInfo,
IsLoaded: isLoaded,
})
}
// QueryCoord builds reopen LoadInfo from DataCoord's full finished-index list for this segment.
// Treat absent index IDs as stale local metadata.
s.fieldIndexes.Range(func(indexID int64, _ *IndexedFieldInfo) bool {
if !indexIDs.Contain(indexID) {
s.fieldIndexes.Remove(indexID)
}
return true
})
}
// Search executes a search on the segment.
// If searchReq.FilterOnly() is true, only executes the filter and returns valid_count (Stage 1 of two-stage search).
func (s *LocalSegment) Search(ctx context.Context, searchReq *segcore.SearchRequest) (*segcore.SearchResult, error) {
filterOnly := searchReq.FilterOnly()
log := mlog.WithLazy(
mlog.Uint64("mvcc", searchReq.MVCC()),
mlog.FieldCollectionID(s.Collection()),
mlog.FieldSegmentID(s.ID()),
mlog.String("segmentType", s.segmentType.String()),
mlog.Bool("filterOnly", filterOnly),
)
if !s.ptrLock.PinIf(state.IsNotReleased) {
// TODO: check if the segment is readable but not released. too many related logic need to be refactor.
return nil, merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
defer s.ptrLock.Unpin()
hasIndex := s.ExistIndex(searchReq.SearchFieldID())
log = log.With(mlog.Bool("withIndex", hasIndex))
log.Debug(ctx, "search segment...")
tr := timerecord.NewTimeRecorder("cgoSearch")
result, err := s.csegment.Search(ctx, searchReq)
if err != nil {
log.Warn(ctx, "Search failed")
return nil, err
}
metrics.QueryNodeSQSegmentLatencyInCore.WithLabelValues(paramtable.GetStringNodeID(), metrics.SearchLabel).Observe(float64(tr.ElapseSpan().Microseconds()) / 1000.0)
if filterOnly {
log.Debug(ctx, "search filter only segment done", mlog.Int64("validCount", result.ValidCount()))
} else {
log.Debug(ctx, "search segment done")
}
return result, nil
}
func (s *LocalSegment) retrieve(ctx context.Context, plan *segcore.RetrievePlan, log *mlog.Logger) (*segcore.RetrieveResult, error) {
if !s.ptrLock.PinIf(state.IsNotReleased) {
// TODO: check if the segment is readable but not released. too many related logic need to be refactor.
return nil, merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
defer s.ptrLock.Unpin()
log.Debug(ctx, "begin to retrieve")
tr := timerecord.NewTimeRecorder("cgoRetrieve")
result, err := s.csegment.Retrieve(ctx, plan)
if err != nil {
log.Warn(ctx, "Retrieve failed")
return nil, err
}
metrics.QueryNodeSQSegmentLatencyInCore.WithLabelValues(paramtable.GetStringNodeID(),
contextutil.GetQueryLabel(ctx)).Observe(float64(tr.ElapseSpan().Microseconds()) / 1000.0)
return result, nil
}
func (s *LocalSegment) Retrieve(ctx context.Context, plan *segcore.RetrievePlan) (*segcorepb.RetrieveResults, error) {
log := mlog.WithLazy(
mlog.FieldCollectionID(s.Collection()),
mlog.FieldPartitionID(s.Partition()),
mlog.FieldSegmentID(s.ID()),
mlog.Uint64("mvcc", plan.Timestamp),
mlog.String("segmentType", s.segmentType.String()),
)
result, err := s.retrieve(ctx, plan, log)
if err != nil {
return nil, err
}
defer result.Release()
_, span := otel.Tracer(typeutil.QueryNodeRole).Start(ctx, "partial-segcore-results-deserialization")
defer span.End()
retrieveResult, err := result.GetResult()
if err != nil {
log.Warn(ctx, "unmarshal retrieve result failed", mlog.Err(err))
return nil, err
}
log.Debug(ctx, "retrieve segment done", mlog.Int("resultNum", len(retrieveResult.Offset)))
return retrieveResult, nil
}
func (s *LocalSegment) retrieveByOffsets(ctx context.Context, plan *segcore.RetrievePlanWithOffsets, log *mlog.Logger) (*segcore.RetrieveResult, error) {
if !s.ptrLock.PinIf(state.IsNotReleased) {
// TODO: check if the segment is readable but not released. too many related logic need to be refactor.
return nil, merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
defer s.ptrLock.Unpin()
log.Debug(ctx, "begin to retrieve by offsets")
tr := timerecord.NewTimeRecorder("cgoRetrieveByOffsets")
result, err := s.csegment.RetrieveByOffsets(ctx, plan)
if err != nil {
log.Warn(ctx, "RetrieveByOffsets failed")
return nil, err
}
metrics.QueryNodeSQSegmentLatencyInCore.WithLabelValues(paramtable.GetStringNodeID(),
contextutil.GetQueryLabel(ctx)).Observe(float64(tr.ElapseSpan().Microseconds()) / 1000.0)
return result, nil
}
func (s *LocalSegment) RetrieveByOffsets(ctx context.Context, plan *segcore.RetrievePlanWithOffsets) (*segcorepb.RetrieveResults, error) {
log := mlog.WithLazy(mlog.FieldCollectionID(s.Collection()),
mlog.FieldPartitionID(s.Partition()),
mlog.FieldSegmentID(s.ID()),
mlog.Int64("msgID", plan.MsgID()),
mlog.String("segmentType", s.segmentType.String()),
mlog.Int("resultNum", len(plan.Offsets)),
)
result, err := s.retrieveByOffsets(ctx, plan, log)
if err != nil {
return nil, err
}
defer result.Release()
_, span := otel.Tracer(typeutil.QueryNodeRole).Start(ctx, "reduced-segcore-results-deserialization")
defer span.End()
retrieveResult, err := result.GetResult()
if err != nil {
log.Warn(ctx, "unmarshal retrieve by offsets result failed", mlog.Err(err))
return nil, err
}
log.Debug(ctx, "retrieve by segment offsets done", mlog.Int("resultNum", len(retrieveResult.Offset)))
return retrieveResult, nil
}
func (s *LocalSegment) Insert(ctx context.Context, rowIDs []int64, timestamps []typeutil.Timestamp, record *segcorepb.InsertRecord) error {
if s.Type() != SegmentTypeGrowing {
return merr.WrapErrServiceInternalMsg("unexpected segmentType when segmentInsert, segmentType = %s", s.segmentType.String())
}
if !s.ptrLock.PinIf(state.IsNotReleased) {
return merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
defer s.ptrLock.Unpin()
var result *segcore.InsertResult
var err error
GetMutatePool().Submit(func() (any, error) {
start := time.Now()
defer func() {
metrics.QueryNodeCGOCallLatency.WithLabelValues(
paramtable.GetStringNodeID(),
"Insert",
"Sync",
).Observe(float64(time.Since(start).Milliseconds()))
}()
result, err = s.csegment.Insert(ctx, &segcore.InsertRequest{
RowIDs: rowIDs,
Timestamps: timestamps,
Record: record,
})
return nil, nil
}).Await()
if err != nil {
return err
}
s.insertCount.Add(result.InsertedRows)
s.rowNum.Store(-1)
s.memSize.Store(-1)
return nil
}
func (s *LocalSegment) Delete(ctx context.Context, primaryKeys storage.PrimaryKeys, timestamps []typeutil.Timestamp) error {
/*
CStatus
Delete(CSegmentInterface c_segment,
long int reserved_offset,
long size,
const long* primary_keys,
const unsigned long* timestamps);
*/
if primaryKeys.Len() == 0 {
return nil
}
if !s.ptrLock.PinIf(state.IsNotReleased) {
return merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
defer s.ptrLock.Unpin()
s.deltaMut.Lock()
defer s.deltaMut.Unlock()
// segcore DeletedRecord::InternalPush is idempotent on (PK, ts):
// duplicate deletes against already-deleted rows are discarded internally.
// Do NOT add a ts-watermark skip here. In L0-forward + partial-L0-compaction
// scenarios, batches contain ts values below the watermark that have NOT yet
// been applied to this segment, and skipping them causes silent data loss.
var err error
GetMutatePool().Submit(func() (any, error) {
start := time.Now()
defer func() {
metrics.QueryNodeCGOCallLatency.WithLabelValues(
paramtable.GetStringNodeID(),
"Delete",
"Sync",
).Observe(float64(time.Since(start).Milliseconds()))
}()
_, err = s.csegment.Delete(ctx, &segcore.DeleteRequest{
PrimaryKeys: primaryKeys,
Timestamps: timestamps,
})
return nil, nil
}).Await()
if err != nil {
return err
}
s.rowNum.Store(-1)
// Track max ts as a high-water-mark (consumed by file-level skip in
// LoadDeltaLogs and by dist_handler for QueryCoord reporting). Using
// tss[last] on unsorted batches underestimates the watermark.
s.advanceLastDeltaTimestamp(timestamps)
return nil
}
// -------------------------------------------------------------------------------------- interfaces for sealed segment
func (s *LocalSegment) LoadFieldData(ctx context.Context, fieldID int64, rowCount int64, field *datapb.FieldBinlog) error {
if !s.ptrLock.PinIf(state.IsNotReleased) {
return merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
defer s.ptrLock.Unpin()
ctx, sp := otel.Tracer(typeutil.QueryNodeRole).Start(ctx, fmt.Sprintf("LoadFieldData-%d-%d", s.ID(), fieldID))
defer sp.End()
log := mlog.With(
mlog.FieldCollectionID(s.Collection()),
mlog.FieldPartitionID(s.Partition()),
mlog.FieldSegmentID(s.ID()),
mlog.FieldFieldID(fieldID),
mlog.Int64("rowCount", rowCount),
)
log.Info(ctx, "start loading field data for field")
// TODO retrieve_enable should be considered
collection := s.collection
fieldSchema, err := getFieldSchema(collection.Schema(), fieldID)
if err != nil {
return err
}
mmapEnabled := isDataMmapEnable(fieldSchema)
fieldWarmupPolicy := getFieldWarmupPolicy(fieldSchema)
req := &segcore.LoadFieldDataRequest{
Fields: []segcore.LoadFieldDataInfo{{
Field: field,
EnableMMap: mmapEnabled,
WarmupPolicy: fieldWarmupPolicy,
}},
RowCount: rowCount,
StorageVersion: s.LoadInfo().GetStorageVersion(),
LoadPriority: s.LoadInfo().GetPriority(),
Shard: s.LoadInfo().GetInsertChannel(),
}
GetLoadPool().Submit(func() (any, error) {
start := time.Now()
defer func() {
metrics.QueryNodeCGOCallLatency.WithLabelValues(
paramtable.GetStringNodeID(),
"LoadFieldData",
"Sync",
).Observe(float64(time.Since(start).Milliseconds()))
}()
_, err = s.csegment.LoadFieldData(ctx, req)
log.Info(ctx, "submitted loadFieldData task to load pool")
return nil, nil
}).Await()
if err != nil {
log.Warn(ctx, "LoadFieldData failed", mlog.Err(err))
return err
}
log.Info(ctx, "load field done")
return nil
}
func (s *LocalSegment) LoadDeltaData(ctx context.Context, deltaData *storage.DeltaData) error {
if deltaData.DeleteRowCount() == 0 {
return nil
}
pks, tss := deltaData.DeletePks(), deltaData.DeleteTimestamps()
rowNum := deltaData.DeleteRowCount()
if !s.ptrLock.PinIf(state.IsNotReleased) {
return merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
defer s.ptrLock.Unpin()
log := mlog.With(
mlog.FieldCollectionID(s.Collection()),
mlog.FieldPartitionID(s.Partition()),
mlog.FieldSegmentID(s.ID()),
)
s.deltaMut.Lock()
defer s.deltaMut.Unlock()
// See comment in Delete(): segcore dedups at (PK, ts) level, and tss is
// NOT sorted across L0 segments (BufferForwarder appends in iteration
// order), so comparing against tss[last] is both unnecessary and incorrect.
ids, err := storage.ParsePrimaryKeysBatch2IDs(pks)
if err != nil {
return err
}
idsBlob, err := proto.Marshal(ids)
if err != nil {
return err
}
loadInfo := C.CLoadDeletedRecordInfo{
timestamps: unsafe.Pointer(&tss[0]),
primary_keys: (*C.uint8_t)(unsafe.Pointer(&idsBlob[0])),
primary_keys_size: C.uint64_t(len(idsBlob)),
row_count: C.int64_t(rowNum),
}
/*
CStatus
LoadDeletedRecord(CSegmentInterface c_segment, CLoadDeletedRecordInfo deleted_record_info)
*/
var status C.CStatus
// Delta-log replay during segment load runs on the load pool, not the
// online-write mutate pool, so a large post-compaction replay cannot starve
// online insert/delete (and thus tSafe advancement).
GetLoadPool().Submit(func() (any, error) {
start := time.Now()
defer func() {
metrics.QueryNodeCGOCallLatency.WithLabelValues(
paramtable.GetStringNodeID(),
"LoadDeletedRecord",
"Sync",
).Observe(float64(time.Since(start).Milliseconds()))
}()
status = C.LoadDeletedRecord(s.ptr, loadInfo)
return nil, nil
}).Await()
if err := HandleCStatus(ctx, &status, "LoadDeletedRecord failed",
mlog.FieldCollectionID(s.Collection()),
mlog.FieldPartitionID(s.Partition()),
mlog.FieldSegmentID(s.ID())); err != nil {
return err
}
s.rowNum.Store(-1)
s.advanceLastDeltaTimestamp(tss)
log.Info(ctx, "load deleted record done",
mlog.Int64("rowNum", rowNum),
mlog.String("segmentType", s.Type().String()))
return nil
}
func GetCLoadInfoWithFunc(ctx context.Context,
fieldSchema *schemapb.FieldSchema,
loadInfo *querypb.SegmentLoadInfo,
indexInfo *querypb.FieldIndexInfo,
f func(c *LoadIndexInfo) error,
) error {
// 1.
loadIndexInfo, err := newLoadIndexInfo(ctx)
if err != nil {
return err
}
defer deleteLoadIndexInfo(loadIndexInfo)
indexParams := funcutil.KeyValuePair2Map(indexInfo.IndexParams)
// as Knowhere reports error if encounter an unknown param, we need to delete it
delete(indexParams, common.MmapEnabledKey)
// 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
}
enableMmap := isIndexMmapEnable(fieldSchema, indexInfo)
// Add warmup policy to index_params if not already present
// C++ will pass it to Knowhere for index loading
if existingWarmup, exists := indexParams[common.WarmupKey]; exists {
mlog.Info(ctx, "warmup policy already in index params (from QueryCoord)",
mlog.FieldSegmentID(loadInfo.GetSegmentID()),
mlog.FieldFieldID(indexInfo.GetFieldID()),
mlog.String("warmup", existingWarmup))
} else {
warmupPolicy := getIndexWarmupPolicy(fieldSchema, indexInfo)
mlog.Info(ctx, "warmup policy from getIndexWarmupPolicy",
mlog.FieldSegmentID(loadInfo.GetSegmentID()),
mlog.FieldFieldID(indexInfo.GetFieldID()),
mlog.String("warmup", warmupPolicy))
if warmupPolicy != "" {
indexParams[common.WarmupKey] = warmupPolicy
}
}
// Pass DataCoord-built index file paths through; QueryNode should not
// attach v0/v1 path layout semantics to the read path.
indexInfoProto := &cgopb.LoadIndexInfo{
CollectionID: loadInfo.GetCollectionID(),
PartitionID: loadInfo.GetPartitionID(),
SegmentID: loadInfo.GetSegmentID(),
Field: fieldSchema,
EnableMmap: enableMmap,
IndexID: indexInfo.GetIndexID(),
IndexBuildID: indexInfo.GetBuildID(),
IndexVersion: indexInfo.GetIndexVersion(),
IndexParams: indexParams,
IndexFiles: indexInfo.GetIndexFilePaths(),
IndexEngineVersion: indexInfo.GetCurrentIndexVersion(),
IndexFileSize: indexInfo.GetIndexSize(),
NumRows: indexInfo.GetNumRows(),
CurrentScalarIndexVersion: indexInfo.GetCurrentScalarIndexVersion(),
IndexStorePathVersion: indexInfo.GetIndexStorePathVersion(),
}
// 2.
if err := loadIndexInfo.appendLoadIndexInfo(ctx, indexInfoProto); err != nil {
mlog.Warn(ctx, "fail to append load index info", mlog.Err(err))
return err
}
loadIndexInfo.setShard(loadInfo.GetInsertChannel())
return f(loadIndexInfo)
}
func (s *LocalSegment) LoadIndex(ctx context.Context, indexInfo *querypb.FieldIndexInfo, fieldType schemapb.DataType) error {
log := mlog.With(
mlog.FieldCollectionID(s.Collection()),
mlog.FieldPartitionID(s.Partition()),
mlog.FieldSegmentID(s.ID()),
mlog.FieldFieldID(indexInfo.GetFieldID()),
mlog.FieldIndexID(indexInfo.GetIndexID()),
)
old := s.GetIndexByID(indexInfo.GetIndexID())
// the index loaded
if old != nil && old.IsLoaded {
log.Warn(ctx, "index already loaded")
return nil
}
ctx, sp := otel.Tracer(typeutil.QueryNodeRole).Start(ctx, fmt.Sprintf("LoadIndex-%d-%d", s.ID(), indexInfo.GetFieldID()))
defer sp.End()
tr := timerecord.NewTimeRecorder("loadIndex")
schemaHelper, err := typeutil.CreateSchemaHelper(s.GetCollection().Schema())
if err != nil {
return err
}
fieldSchema, err := schemaHelper.GetFieldFromID(indexInfo.GetFieldID())
if err != nil {
return err
}
// // if segment is pk sorted, user created indexes bring no performance gain but extra memory usage
if s.IsSorted() && fieldSchema.GetIsPrimaryKey() {
log.Info(ctx, "skip loading index for pk field in sorted segment")
// set field index, preventing repeated loading index task
s.fieldIndexes.Insert(indexInfo.GetFieldID(), &IndexedFieldInfo{
FieldBinlog: &datapb.FieldBinlog{
FieldID: indexInfo.GetFieldID(),
},
IndexInfo: indexInfo,
IsLoaded: true,
})
return nil
}
return s.innerLoadIndex(ctx, fieldSchema, indexInfo, tr, fieldType)
}
func (s *LocalSegment) innerLoadIndex(ctx context.Context,
fieldSchema *schemapb.FieldSchema,
indexInfo *querypb.FieldIndexInfo,
tr *timerecord.TimeRecorder,
fieldType schemapb.DataType,
) error {
err := GetCLoadInfoWithFunc(ctx, fieldSchema,
s.LoadInfo(), indexInfo, func(loadIndexInfo *LoadIndexInfo) error {
newLoadIndexInfoSpan := tr.RecordSpan()
if err := loadIndexInfo.loadIndex(ctx); err != nil {
if loadIndexInfo.cleanLocalData(ctx) != nil {
mlog.Warn(ctx, "failed to clean cached data on disk after append index failed",
mlog.FieldBuildID(indexInfo.BuildID),
mlog.Int64("index version", indexInfo.IndexVersion))
}
return err
}
if s.Type() != SegmentTypeSealed {
return merr.WrapErrServiceInternalMsg("updateSegmentIndex failed, illegal segment type %s, segmentID = %d", s.segmentType, s.ID())
}
appendLoadIndexInfoSpan := tr.RecordSpan()
// 3.
err := s.UpdateIndexInfo(ctx, indexInfo, loadIndexInfo)
if err != nil {
return err
}
updateIndexInfoSpan := tr.RecordSpan()
mlog.Info(ctx, "Finish loading index",
mlog.Duration("newLoadIndexInfoSpan", newLoadIndexInfoSpan),
mlog.Duration("appendLoadIndexInfoSpan", appendLoadIndexInfoSpan),
mlog.Duration("updateIndexInfoSpan", updateIndexInfoSpan),
)
return nil
})
if err != nil {
mlog.Warn(ctx, "load index failed", mlog.Err(err))
}
return err
}
func (s *LocalSegment) UpdateIndexInfo(ctx context.Context, indexInfo *querypb.FieldIndexInfo, info *LoadIndexInfo) error {
log := mlog.With(
mlog.FieldCollectionID(s.Collection()),
mlog.FieldPartitionID(s.Partition()),
mlog.FieldSegmentID(s.ID()),
mlog.FieldFieldID(indexInfo.FieldID),
)
if !s.ptrLock.PinIf(state.IsNotReleased) {
return merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
defer s.ptrLock.Unpin()
var status C.CStatus
GetDynamicPool().Submit(func() (any, error) {
status = C.UpdateSealedSegmentIndex(s.ptr, info.cLoadIndexInfo)
return nil, nil
}).Await()
if err := HandleCStatus(ctx, &status, "UpdateSealedSegmentIndex failed",
mlog.FieldCollectionID(s.Collection()),
mlog.FieldPartitionID(s.Partition()),
mlog.FieldSegmentID(s.ID()),
mlog.FieldFieldID(indexInfo.FieldID)); err != nil {
return err
}
s.fieldIndexes.Insert(indexInfo.GetIndexID(), &IndexedFieldInfo{
FieldBinlog: &datapb.FieldBinlog{
FieldID: indexInfo.GetFieldID(),
},
IndexInfo: indexInfo,
IsLoaded: true,
})
log.Info(ctx, "updateSegmentIndex done")
return nil
}
func (s *LocalSegment) UpdateFieldRawDataSize(ctx context.Context, numRows int64, fieldBinlog *datapb.FieldBinlog) error {
var status C.CStatus
fieldID := fieldBinlog.FieldID
fieldDataSize := int64(0)
for _, binlog := range fieldBinlog.GetBinlogs() {
fieldDataSize += binlog.GetMemorySize()
}
GetDynamicPool().Submit(func() (any, error) {
status = C.UpdateFieldRawDataSize(s.ptr, C.int64_t(fieldID), C.int64_t(numRows), C.int64_t(fieldDataSize))
return nil, nil
}).Await()
if err := HandleCStatus(ctx, &status, "updateFieldRawDataSize failed"); err != nil {
return err
}
mlog.Info(ctx, "updateFieldRawDataSize done", mlog.FieldSegmentID(s.ID()))
return nil
}
func (s *LocalSegment) syncFieldJSONStatsFromLoadInfo(ctx context.Context, loadInfo *querypb.SegmentLoadInfo) {
jsonStatsInfo := make(map[int64]*querypb.JsonStatsInfo)
if !paramtable.Get().CommonCfg.EnabledJSONKeyStats.GetAsBool() {
mlog.Warn(ctx, "skip sync json key stats, json key stats is not enabled", mlog.FieldSegmentID(s.ID()))
s.fieldJSONStatsMu.Lock()
s.fieldJSONStats = jsonStatsInfo
s.fieldJSONStatsMu.Unlock()
return
}
statsResult := packed.NewStatsResolverFromLoadInfo(loadInfo).TextAndJSONIndexStatsWithBasePaths()
jsonKeyStats := statsResult.JSONKeyStats
if statsResult.Err() != nil {
mlog.Warn(ctx, "failed to resolve json key stats from manifest",
mlog.FieldSegmentID(loadInfo.GetSegmentID()),
mlog.String("manifestPath", loadInfo.GetManifestPath()),
mlog.Err(statsResult.Err()))
jsonKeyStats = loadInfo.GetJsonKeyStatsLogs()
}
for fieldID, stats := range jsonKeyStats {
if stats == nil {
continue
}
if stats.GetJsonKeyStatsDataFormat() != common.JSONStatsDataFormatVersion {
mlog.Warn(ctx, "skip sync json key stats, data format invalid",
mlog.FieldSegmentID(loadInfo.GetSegmentID()),
mlog.FieldFieldID(fieldID),
mlog.FieldBuildID(stats.GetBuildID()),
mlog.Int64("version", stats.GetVersion()),
mlog.Int64("dataFormat", stats.GetJsonKeyStatsDataFormat()),
mlog.Int64("expectedDataFormat", common.JSONStatsDataFormatVersion))
continue
}
jsonStatsInfo[fieldID] = &querypb.JsonStatsInfo{
FieldID: stats.GetFieldID(),
DataFormatVersion: stats.GetJsonKeyStatsDataFormat(),
BuildID: stats.GetBuildID(),
VersionID: stats.GetVersion(),
}
}
s.fieldJSONStatsMu.Lock()
s.fieldJSONStats = jsonStatsInfo
s.fieldJSONStatsMu.Unlock()
}
func (s *LocalSegment) Load(ctx context.Context) error {
if err := s.csegment.Load(ctx); err != nil {
return err
}
s.syncFieldJSONStatsFromLoadInfo(ctx, s.LoadInfo())
return nil
}
func (s *LocalSegment) Reopen(ctx context.Context, newLoadInfo *querypb.SegmentLoadInfo) error {
if !s.ptrLock.PinIfNotReleased() {
return merr.WrapErrSegmentNotLoaded(s.ID(), "segment released during reopen")
}
defer s.ptrLock.Unpin()
schema, schemaVersion := s.collection.SchemaAndSegcoreVersion()
err := s.csegment.Reopen(ctx, &segcore.ReopenRequest{
LoadInfo: newLoadInfo,
Schema: schema,
SchemaVersion: schemaVersion,
})
if err != nil {
return err
}
s.syncFieldIndexes(newLoadInfo.GetIndexInfos())
s.loadInfo.Store(newLoadInfo)
s.syncFieldJSONStatsFromLoadInfo(ctx, newLoadInfo)
return nil
}
type ReleaseScope int
const (
ReleaseScopeAll ReleaseScope = iota
ReleaseScopeData
)
type releaseOptions struct {
Scope ReleaseScope
}
func newReleaseOptions() *releaseOptions {
return &releaseOptions{
Scope: ReleaseScopeAll,
}
}
type releaseOption func(*releaseOptions)
func WithReleaseScope(scope ReleaseScope) releaseOption {
return func(options *releaseOptions) {
options.Scope = scope
}
}
func (s *LocalSegment) Release(ctx context.Context, opts ...releaseOption) {
options := newReleaseOptions()
for _, opt := range opts {
opt(options)
}
stateLockGuard := s.startRelease(options.Scope)
if stateLockGuard == nil { // release is already done.
return
}
// release will never fail
defer stateLockGuard.Done(nil)
log := mlog.With(mlog.FieldCollectionID(s.Collection()),
mlog.FieldPartitionID(s.Partition()),
mlog.FieldSegmentID(s.ID()),
mlog.String("segmentType", s.segmentType.String()),
mlog.Int64("insertCount", s.InsertCount()),
)
// wait all read ops finished
ptr := s.ptr
if options.Scope == ReleaseScopeData {
s.ReleaseSegmentData()
log.Info(ctx, "release segment data done and the field indexes info has been set lazy load=true")
return
}
if paramtable.Get().QueryNodeCfg.ExprResCacheEnabled.GetAsBool() {
// erase expr-cache for this segment before deleting C segment
C.ExprResCacheEraseSegment(C.int64_t(s.ID()))
}
GetDynamicPool().Submit(func() (any, error) {
C.DeleteSegment(ptr)
return nil, nil
}).Await()
// TODO: disable logical resource handling for now
// usage := s.ResourceUsageEstimate()
// s.manager.SubLogicalResource(usage)
// Refund PK candidate resource
s.pkCandidate.Refund()
binlogSize := s.binlogSize.Load()
if binlogSize > 0 {
// no concurrent change to s.binlogSize, so the subtraction is safe
s.manager.SubLoadedBinlogSize(binlogSize)
s.binlogSize.Store(0)
}
log.Info(ctx, "delete segment from memory")
}
// ReleaseSegmentData releases the segment data.
func (s *LocalSegment) ReleaseSegmentData() {
GetDynamicPool().Submit(func() (any, error) {
C.ClearSegmentData(s.ptr)
return nil, nil
}).Await()
for _, indexInfo := range s.Indexes() {
indexInfo.IsLoaded = false
}
}
// StartLoadData starts the loading process of the segment.
func (s *LocalSegment) StartLoadData() (state.LoadStateLockGuard, error) {
return s.ptrLock.StartLoadData()
}
// startRelease starts the releasing process of the segment.
func (s *LocalSegment) startRelease(scope ReleaseScope) state.LoadStateLockGuard {
switch scope {
case ReleaseScopeData:
return s.ptrLock.StartReleaseData()
case ReleaseScopeAll:
return s.ptrLock.StartReleaseAll()
default:
panic(fmt.Sprintf("unexpected release scope %d", scope))
}
}
func (s *LocalSegment) RemoveFieldFile(fieldId int64) {
GetDynamicPool().Submit(func() (any, error) {
C.RemoveFieldFile(s.ptr, C.int64_t(fieldId))
return nil, nil
}).Await()
}
func (s *LocalSegment) RemoveUnusedFieldFiles() error {
schema := s.collection.Schema()
indexInfos, _ := separateIndexAndBinlog(s.LoadInfo())
for _, indexInfo := range indexInfos {
need, err := s.indexNeedLoadRawData(schema, indexInfo)
if err != nil {
return err
}
if !need {
s.RemoveFieldFile(indexInfo.IndexInfo.FieldID)
}
}
return nil
}
func (s *LocalSegment) indexNeedLoadRawData(schema *schemapb.CollectionSchema, indexInfo *IndexedFieldInfo) (bool, error) {
schemaHelper, err := typeutil.CreateSchemaHelper(schema)
if err != nil {
return false, err
}
fieldSchema, err := schemaHelper.GetFieldFromID(indexInfo.IndexInfo.FieldID)
if err != nil {
return false, err
}
return !typeutil.IsVectorType(fieldSchema.DataType) && s.HasRawData(indexInfo.IndexInfo.FieldID), nil
}
func (s *LocalSegment) GetFieldJSONIndexStats() map[int64]*querypb.JsonStatsInfo {
s.fieldJSONStatsMu.RLock()
defer s.fieldJSONStatsMu.RUnlock()
stats := make(map[int64]*querypb.JsonStatsInfo, len(s.fieldJSONStats))
for fieldID, info := range s.fieldJSONStats {
if info == nil {
continue
}
stats[fieldID] = proto.Clone(info).(*querypb.JsonStatsInfo)
}
return stats
}
// MaterializedFieldIDs returns the field ids with materialized columns in the
// growing segment's insert record. The flush layout must be trimmed to this
// set; a non-materialized column is legally absent (a dropped field or a
// function output backfilled by bump-schema compaction).
func (s *LocalSegment) MaterializedFieldIDs(ctx context.Context) ([]int64, error) {
if s.Type() != SegmentTypeGrowing {
return nil, merr.WrapErrServiceInternalMsg("unexpected segmentType for MaterializedFieldIDs, segmentType = %s", s.segmentType.String())
}
if !s.ptrLock.PinIf(state.IsNotReleased) {
return nil, merr.WrapErrSegmentNotLoaded(s.ID(), "segment released")
}
defer s.ptrLock.Unpin()
var cIDs *C.int64_t
var cCount C.int64_t
var status C.CStatus
GetDynamicPool().Submit(func() (any, error) {
status = C.GetGrowingSegmentMaterializedFieldIDs(s.ptr, &cIDs, &cCount)
return nil, nil
}).Await()
if err := HandleCStatus(ctx, &status, "GetGrowingSegmentMaterializedFieldIDs"); err != nil {
return nil, err
}
if cIDs == nil || cCount == 0 {
return nil, nil
}
defer C.free(unsafe.Pointer(cIDs))
ids := make([]int64, 0, int(cCount))
for _, id := range unsafe.Slice(cIDs, int(cCount)) {
ids = append(ids, int64(id))
}
return ids, nil
}
// FlushData flushes data from the growing segment directly to storage via C++ milvus-storage.
// This is a unified interface that combines data extraction from segcore and writing to storage.
// The C++ side handles: extracting raw field data from ConcurrentVector, converting to Arrow,
// and writing to storage via milvus-storage with TEXT column LOB handling.
func (s *LocalSegment) FlushData(ctx context.Context, startOffset, endOffset int64, config *FlushConfig) (*FlushResult, error) {
// currently only growing segments support FlushData
if s.Type() != SegmentTypeGrowing {
return nil, merr.WrapErrServiceInternalMsg("FlushData is only supported for growing segments, got %s", s.Type().String())
}
// validate offsets
if startOffset < 0 || endOffset < startOffset {
return nil, merr.WrapErrServiceInternalMsg("invalid offsets: start=%d, end=%d", startOffset, endOffset)
}
// no data to flush
if startOffset == endOffset {
return nil, nil
}
if config == nil {
return nil, merr.WrapErrServiceInternalMsg("flush config is nil")
}
if config.Schema == nil {
return nil, merr.WrapErrServiceInternalMsg("flush schema is nil")
}
schemaBlob, err := proto.Marshal(config.Schema)
if err != nil {
return nil, merr.WrapErrServiceInternalMsg("marshal flush schema failed: %s", err.Error())
}
if len(schemaBlob) == 0 {
return nil, merr.WrapErrServiceInternalMsg("marshal flush schema returned empty blob")
}
// build C flush config
var cConfig C.CFlushConfig
cSegmentPath := C.CString(config.SegmentBasePath)
defer C.free(unsafe.Pointer(cSegmentPath))
cConfig.segment_path = cSegmentPath
cSchemaBlob := C.CBytes(schemaBlob)
defer C.free(cSchemaBlob)
cConfig.schema_blob = cSchemaBlob
cConfig.schema_length = C.int64_t(len(schemaBlob))
cConfig.read_version = C.int64_t(config.ReadVersion)
cConfig.retry_limit = C.uint32_t(3)
writerFormat := config.WriterFormat
if writerFormat != "" {
cWriterFormat := C.CString(writerFormat)
defer C.free(unsafe.Pointer(cWriterFormat))
cConfig.writer_format = cWriterFormat
}
schemaBasedPattern := config.SchemaBasedPattern
if schemaBasedPattern != "" {
cSchemaBasedPattern := C.CString(schemaBasedPattern)
defer C.free(unsafe.Pointer(cSchemaBasedPattern))
cConfig.schema_based_pattern = cSchemaBasedPattern
}
schemaBasedFormats := config.SchemaBasedFormats
if schemaBasedFormats != "" {
cSchemaBasedFormats := C.CString(schemaBasedFormats)
defer C.free(unsafe.Pointer(cSchemaBasedFormats))
cConfig.schema_based_formats = cSchemaBasedFormats
}
numAllowedFields := len(config.AllowedFieldIDs)
if numAllowedFields > 0 {
cAllowedFieldIDs := (*C.int64_t)(C.malloc(C.size_t(numAllowedFields) * C.size_t(unsafe.Sizeof(C.int64_t(0)))))
defer C.free(unsafe.Pointer(cAllowedFieldIDs))
allowedFieldIDSlice := unsafe.Slice(cAllowedFieldIDs, numAllowedFields)
for i, fieldID := range config.AllowedFieldIDs {
allowedFieldIDSlice[i] = C.int64_t(fieldID)
}
cConfig.allowed_field_ids = cAllowedFieldIDs
cConfig.num_allowed_fields = C.size_t(numAllowedFields)
}
numColumnGroups := len(config.ColumnGroups)
if numColumnGroups > 0 {
totalFieldCount := 0
for _, columnGroup := range config.ColumnGroups {
totalFieldCount += len(columnGroup.Fields)
}
cColumnGroupIDs := (*C.int64_t)(C.malloc(C.size_t(numColumnGroups) * C.size_t(unsafe.Sizeof(C.int64_t(0)))))
defer C.free(unsafe.Pointer(cColumnGroupIDs))
cColumnGroupFieldCounts := (*C.size_t)(C.malloc(C.size_t(numColumnGroups) * C.size_t(unsafe.Sizeof(C.size_t(0)))))
defer C.free(unsafe.Pointer(cColumnGroupFieldCounts))
var cColumnGroupFieldIDs *C.int64_t
if totalFieldCount > 0 {
cColumnGroupFieldIDs = (*C.int64_t)(C.malloc(C.size_t(totalFieldCount) * C.size_t(unsafe.Sizeof(C.int64_t(0)))))
defer C.free(unsafe.Pointer(cColumnGroupFieldIDs))
}
groupIDSlice := unsafe.Slice(cColumnGroupIDs, numColumnGroups)
fieldCountSlice := unsafe.Slice(cColumnGroupFieldCounts, numColumnGroups)
var fieldIDSlice []C.int64_t
if totalFieldCount > 0 {
fieldIDSlice = unsafe.Slice(cColumnGroupFieldIDs, totalFieldCount)
}
fieldOffset := 0
for i, columnGroup := range config.ColumnGroups {
groupIDSlice[i] = C.int64_t(columnGroup.GroupID)
fieldCountSlice[i] = C.size_t(len(columnGroup.Fields))
for _, fieldID := range columnGroup.Fields {
fieldIDSlice[fieldOffset] = C.int64_t(fieldID)
fieldOffset++
}
}
cConfig.column_group_ids = cColumnGroupIDs
cConfig.column_group_field_ids = cColumnGroupFieldIDs
cConfig.column_group_field_counts = cColumnGroupFieldCounts
cConfig.num_column_groups = C.size_t(numColumnGroups)
}
// populate TEXT column configs
// All arrays must be C-allocated to avoid "Go pointer to unpinned Go pointer" panic
numTextCols := len(config.TextFieldIDs)
if numTextCols > 0 {
// allocate C arrays via C.malloc (not Go slices) to satisfy CGO pointer rules
cFieldIDs := (*C.int64_t)(C.malloc(C.size_t(numTextCols) * C.size_t(unsafe.Sizeof(C.int64_t(0)))))
defer C.free(unsafe.Pointer(cFieldIDs))
cLobPathPtrs := (**C.char)(C.malloc(C.size_t(numTextCols) * C.size_t(unsafe.Sizeof((*C.char)(nil)))))
defer C.free(unsafe.Pointer(cLobPathPtrs))
fieldIDSlice := unsafe.Slice(cFieldIDs, numTextCols)
lobPathSlice := unsafe.Slice(cLobPathPtrs, numTextCols)
for i := 0; i < numTextCols; i++ {
fieldIDSlice[i] = C.int64_t(config.TextFieldIDs[i])
lobPathSlice[i] = C.CString(config.TextLobPaths[i])
defer C.free(unsafe.Pointer(lobPathSlice[i]))
}
cConfig.text_field_ids = cFieldIDs
cConfig.text_lob_paths = cLobPathPtrs
cConfig.text_inline_threshold = C.int64_t(config.TextInlineThreshold)
cConfig.text_max_lob_file_bytes = C.int64_t(config.TextMaxLobFileBytes)
cConfig.text_flush_threshold_bytes = C.int64_t(config.TextFlushThresholdBytes)
cConfig.num_text_columns = C.size_t(numTextCols)
} else {
cConfig.text_field_ids = nil
cConfig.text_lob_paths = nil
cConfig.num_text_columns = 0
}
numBM25Fields := len(config.BM25FieldIDs)
if numBM25Fields > 0 {
if len(config.BM25StatsLogIDs) != numBM25Fields {
return nil, merr.WrapErrServiceInternalMsg("BM25 stats log IDs count mismatch, fields=%d logIDs=%d", numBM25Fields, len(config.BM25StatsLogIDs))
}
cBM25FieldIDs := (*C.int64_t)(C.malloc(C.size_t(numBM25Fields) * C.size_t(unsafe.Sizeof(C.int64_t(0)))))
defer C.free(unsafe.Pointer(cBM25FieldIDs))
cBM25StatsLogIDs := (*C.int64_t)(C.malloc(C.size_t(numBM25Fields) * C.size_t(unsafe.Sizeof(C.int64_t(0)))))
defer C.free(unsafe.Pointer(cBM25StatsLogIDs))
bm25FieldIDSlice := unsafe.Slice(cBM25FieldIDs, numBM25Fields)
bm25StatsLogIDSlice := unsafe.Slice(cBM25StatsLogIDs, numBM25Fields)
for i, fieldID := range config.BM25FieldIDs {
bm25FieldIDSlice[i] = C.int64_t(fieldID)
bm25StatsLogIDSlice[i] = C.int64_t(config.BM25StatsLogIDs[i])
}
cConfig.bm25_field_ids = cBM25FieldIDs
cConfig.bm25_stats_log_ids = cBM25StatsLogIDs
cConfig.num_bm25_fields = C.size_t(numBM25Fields)
} else {
cConfig.bm25_field_ids = nil
cConfig.bm25_stats_log_ids = nil
cConfig.num_bm25_fields = 0
}
cConfig.write_merged_bm25_stats = C.bool(config.WriteMergedBM25Stats)
// call C FFI
var cResult C.CFlushResult
status := C.FlushGrowingSegmentData(
s.ptr,
C.int64_t(startOffset),
C.int64_t(endOffset),
&cConfig,
&cResult,
)
defer C.FreeFlushResult(&cResult)
if err := HandleCStatus(ctx, &status, "FlushGrowingSegmentData"); err != nil {
return nil, err
}
// no data flushed
if cResult.manifest_path == nil {
return nil, nil
}
// C++ SegmentWriter returns a raw file path like:
// "files/insert_log/{coll}/{part}/{seg}/_metadata/manifest-{ver}.avro"
// But GetLoonManifest() expects a JSON-encoded manifest path like:
// {"ver":{ver},"base_path":"files/insert_log/{coll}/{part}/{seg}"}
// Convert using the committed_version and base path extraction.
rawPath := C.GoString(cResult.manifest_path)
committedVersion := int64(cResult.committed_version)
// Extract base path: strip "/_metadata/manifest-{ver}.avro" suffix
basePath := rawPath
if idx := strings.Index(rawPath, "/_metadata/"); idx >= 0 {
basePath = rawPath[:idx]
}
manifestPath := packed.MarshalManifestPath(basePath, committedVersion)
fieldNullCounts := make(map[int64]int64, int(cResult.num_field_stats))
if cResult.num_field_stats > 0 {
fieldIDs := unsafe.Slice(cResult.field_ids, int(cResult.num_field_stats))
nullCounts := unsafe.Slice(cResult.field_null_counts, int(cResult.num_field_stats))
for i := 0; i < int(cResult.num_field_stats); i++ {
fieldID := int64(fieldIDs[i])
fieldNullCounts[fieldID] = int64(nullCounts[i])
}
}
flushedFieldIDs := make([]int64, 0, int(cResult.num_flushed_fields))
if cResult.num_flushed_fields > 0 {
ids := unsafe.Slice(cResult.flushed_field_ids, int(cResult.num_flushed_fields))
for i := 0; i < int(cResult.num_flushed_fields); i++ {
flushedFieldIDs = append(flushedFieldIDs, int64(ids[i]))
}
}
columnGroupMemorySizes := make(map[int64]int64, int(cResult.num_column_groups))
if cResult.num_column_groups > 0 {
columnGroupIDs := unsafe.Slice(cResult.column_group_ids, int(cResult.num_column_groups))
memorySizes := unsafe.Slice(cResult.column_group_memory_sizes, int(cResult.num_column_groups))
for i := 0; i < int(cResult.num_column_groups); i++ {
columnGroupMemorySizes[int64(columnGroupIDs[i])] = int64(memorySizes[i])
}
}
bm25Stats := make(map[int64]*storage.BM25Stats, int(cResult.num_bm25_stats))
if cResult.num_bm25_stats > 0 {
fieldIDs := unsafe.Slice(cResult.bm25_field_ids, int(cResult.num_bm25_stats))
statsBytes := unsafe.Slice(cResult.bm25_stats, int(cResult.num_bm25_stats))
statsSizes := unsafe.Slice(cResult.bm25_stats_sizes, int(cResult.num_bm25_stats))
for i := 0; i < int(cResult.num_bm25_stats); i++ {
bytes := C.GoBytes(unsafe.Pointer(statsBytes[i]), C.int(statsSizes[i]))
stats, err := storage.NewBM25StatsWithBytes(bytes)
if err != nil {
return nil, err
}
bm25Stats[int64(fieldIDs[i])] = stats
}
}
return &FlushResult{
ManifestPath: manifestPath,
NumRows: int64(cResult.num_rows),
TimestampFrom: uint64(cResult.timestamp_from),
TimestampTo: uint64(cResult.timestamp_to),
FlushedFieldIDs: flushedFieldIDs,
ColumnGroupMemorySizes: columnGroupMemorySizes,
FieldNullCounts: fieldNullCounts,
BM25Stats: bm25Stats,
}, nil
}