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
579 lines
20 KiB
Go
579 lines
20 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 querynodev2
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strconv"
|
|
|
|
"github.com/samber/lo"
|
|
"go.opentelemetry.io/otel/trace"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
|
"github.com/milvus-io/milvus/internal/querynodev2/delegator"
|
|
"github.com/milvus-io/milvus/internal/querynodev2/segments"
|
|
"github.com/milvus-io/milvus/internal/querynodev2/tasks"
|
|
"github.com/milvus-io/milvus/internal/util/reduce"
|
|
"github.com/milvus-io/milvus/internal/util/segmentutil"
|
|
"github.com/milvus-io/milvus/internal/util/streamrpc"
|
|
"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/internalpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
|
|
"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/merr"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
|
|
)
|
|
|
|
func loadL0Segments(ctx context.Context, delegator delegator.ShardDelegator, req *querypb.WatchDmChannelsRequest) error {
|
|
l0Segments := make([]*querypb.SegmentLoadInfo, 0)
|
|
for _, channel := range req.GetInfos() {
|
|
for _, segmentID := range channel.GetLevelZeroSegmentIds() {
|
|
segmentInfo, ok := req.GetSegmentInfos()[segmentID]
|
|
if !ok ||
|
|
segmentInfo.GetLevel() != datapb.SegmentLevel_L0 {
|
|
continue
|
|
}
|
|
|
|
l0Segments = append(l0Segments, &querypb.SegmentLoadInfo{
|
|
SegmentID: segmentID,
|
|
PartitionID: segmentInfo.PartitionID,
|
|
CollectionID: segmentInfo.CollectionID,
|
|
BinlogPaths: segmentInfo.Binlogs,
|
|
NumOfRows: segmentInfo.NumOfRows,
|
|
Statslogs: segmentInfo.Statslogs,
|
|
Deltalogs: segmentInfo.Deltalogs,
|
|
Bm25Logs: segmentInfo.Bm25Statslogs,
|
|
InsertChannel: segmentInfo.InsertChannel,
|
|
StartPosition: segmentInfo.GetStartPosition(),
|
|
Level: segmentInfo.GetLevel(),
|
|
})
|
|
}
|
|
}
|
|
|
|
return delegator.LoadL0(ctx, l0Segments, req.GetVersion())
|
|
}
|
|
|
|
func loadGrowingSegments(ctx context.Context, delegator delegator.ShardDelegator, req *querypb.WatchDmChannelsRequest) error {
|
|
// load growing segments
|
|
growingSegments := make([]*querypb.SegmentLoadInfo, 0, len(req.Infos))
|
|
for _, info := range req.Infos {
|
|
for _, segmentID := range info.GetUnflushedSegmentIds() {
|
|
// unFlushed segment may not have binLogs, skip loading
|
|
segmentInfo := req.GetSegmentInfos()[segmentID]
|
|
if segmentInfo == nil {
|
|
mlog.Warn(ctx, "an unflushed segment is not found in segment infos", mlog.FieldSegmentID(segmentID))
|
|
continue
|
|
}
|
|
if len(segmentInfo.GetBinlogs()) > 0 || segmentInfo.GetManifestPath() != "" {
|
|
growingSegments = append(growingSegments, &querypb.SegmentLoadInfo{
|
|
SegmentID: segmentInfo.ID,
|
|
Level: segmentInfo.GetLevel(),
|
|
PartitionID: segmentInfo.PartitionID,
|
|
CollectionID: segmentInfo.CollectionID,
|
|
BinlogPaths: segmentInfo.Binlogs,
|
|
NumOfRows: segmentInfo.NumOfRows,
|
|
Statslogs: segmentInfo.Statslogs,
|
|
Deltalogs: segmentInfo.Deltalogs,
|
|
Bm25Logs: segmentInfo.Bm25Statslogs,
|
|
InsertChannel: segmentInfo.InsertChannel,
|
|
StartPosition: segmentInfo.GetStartPosition(),
|
|
StorageVersion: segmentInfo.GetStorageVersion(),
|
|
ManifestPath: segmentInfo.GetManifestPath(),
|
|
})
|
|
} else {
|
|
mlog.Info(ctx, "skip segment which has no binlog and no manifest path",
|
|
mlog.FieldSegmentID(segmentInfo.ID),
|
|
mlog.String("manifestPath", segmentInfo.GetManifestPath()))
|
|
}
|
|
}
|
|
}
|
|
|
|
return delegator.LoadGrowing(ctx, growingSegments, req.GetVersion())
|
|
}
|
|
|
|
func (node *QueryNode) loadDeltaLogs(ctx context.Context, req *querypb.LoadSegmentsRequest) *commonpb.Status {
|
|
log := mlog.With(
|
|
mlog.FieldCollectionID(req.GetCollectionID()),
|
|
)
|
|
|
|
var finalErr error
|
|
for _, info := range req.GetInfos() {
|
|
segment := node.manager.Segment.GetSealed(info.GetSegmentID())
|
|
if segment == nil {
|
|
continue
|
|
}
|
|
|
|
err := node.loader.LoadDeltaLogs(ctx, segment, info)
|
|
if err != nil {
|
|
if finalErr == nil {
|
|
finalErr = err
|
|
}
|
|
continue
|
|
}
|
|
// try to update segment version after load delta logs
|
|
node.manager.Segment.UpdateBy(segments.IncreaseVersion(req.GetVersion()), segments.WithType(segments.SegmentTypeSealed), segments.WithID(info.GetSegmentID()))
|
|
}
|
|
|
|
if finalErr != nil {
|
|
log.Warn(ctx, "failed to load delta logs", mlog.Err(finalErr))
|
|
return merr.Status(finalErr)
|
|
}
|
|
|
|
return merr.Success()
|
|
}
|
|
|
|
func (node *QueryNode) loadIndex(ctx context.Context, req *querypb.LoadSegmentsRequest) *commonpb.Status {
|
|
log := mlog.With(
|
|
mlog.FieldCollectionID(req.GetCollectionID()),
|
|
mlog.Int64s("segmentIDs", lo.Map(req.GetInfos(), func(info *querypb.SegmentLoadInfo, _ int) int64 { return info.GetSegmentID() })),
|
|
)
|
|
|
|
status := merr.Success()
|
|
log.Info(ctx, "start to load index")
|
|
|
|
for _, info := range req.GetInfos() {
|
|
log := mlog.With(mlog.FieldSegmentID(info.GetSegmentID()))
|
|
segment := node.manager.Segment.GetSealed(info.GetSegmentID())
|
|
if segment == nil {
|
|
log.Warn(ctx, "segment not found for load index operation")
|
|
continue
|
|
}
|
|
localSegment, ok := segment.(*segments.LocalSegment)
|
|
if !ok {
|
|
log.Warn(ctx, "segment not local for load index opeartion")
|
|
continue
|
|
}
|
|
|
|
err := node.loader.LoadIndex(ctx, localSegment, info, req.Version)
|
|
if err != nil {
|
|
log.Warn(ctx, "failed to load index", mlog.Err(err))
|
|
status = merr.Status(err)
|
|
break
|
|
}
|
|
}
|
|
|
|
return status
|
|
}
|
|
|
|
func (node *QueryNode) reopenSegments(ctx context.Context, req *querypb.LoadSegmentsRequest) *commonpb.Status {
|
|
log := mlog.With(
|
|
mlog.FieldCollectionID(req.GetCollectionID()),
|
|
mlog.Int64s("segmentIDs", lo.Map(req.GetInfos(), func(info *querypb.SegmentLoadInfo, _ int) int64 { return info.GetSegmentID() })),
|
|
)
|
|
|
|
log.Info(ctx, "start to reopen segments")
|
|
err := node.loader.ReopenSegments(ctx, req.GetInfos())
|
|
if err != nil {
|
|
log.Warn(ctx, "failed to reopen segments", mlog.Err(err))
|
|
return merr.Status(err)
|
|
}
|
|
return merr.Success()
|
|
}
|
|
|
|
func (node *QueryNode) queryChannel(ctx context.Context, req *querypb.QueryRequest, channel string) (*internalpb.RetrieveResults, error) {
|
|
queryLabel := req.GetReq().GetQueryLabel()
|
|
if queryLabel == "" {
|
|
queryLabel = metrics.QueryLabel
|
|
}
|
|
ctx = contextutil.WithQueryLabel(ctx, queryLabel)
|
|
msgID := req.Req.Base.GetMsgID()
|
|
traceID := trace.SpanFromContext(ctx).SpanContext().TraceID()
|
|
log := mlog.With(
|
|
mlog.Int64("msgID", msgID),
|
|
mlog.FieldCollectionID(req.GetReq().GetCollectionID()),
|
|
mlog.String("channel", channel),
|
|
mlog.String("scope", req.GetScope().String()),
|
|
mlog.String("queryLabel", queryLabel),
|
|
)
|
|
|
|
var err error
|
|
metrics.QueryNodeSQCount.WithLabelValues(fmt.Sprint(node.GetNodeID()), queryLabel, metrics.TotalLabel, metrics.Leader, fmt.Sprint(req.GetReq().GetCollectionID())).Inc()
|
|
defer func() {
|
|
if err != nil {
|
|
metrics.QueryNodeSQCount.WithLabelValues(fmt.Sprint(node.GetNodeID()), queryLabel, metrics.FailLabel, metrics.Leader, fmt.Sprint(req.GetReq().GetCollectionID())).Inc()
|
|
}
|
|
metrics.QueryNodePartialResultCount.WithLabelValues(fmt.Sprint(node.GetNodeID()), queryLabel, fmt.Sprint(req.GetReq().GetCollectionID())).Inc()
|
|
}()
|
|
|
|
log.Debug(ctx, "start do query with channel",
|
|
mlog.Int64s("segmentIDs", req.GetSegmentIDs()),
|
|
)
|
|
// add cancel when error occurs
|
|
queryCtx, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
|
|
// From Proxy
|
|
tr := timerecord.NewTimeRecorder("queryDelegator")
|
|
// get delegator
|
|
sd, ok := node.delegators.Get(channel)
|
|
if !ok {
|
|
err := merr.WrapErrChannelNotFound(channel)
|
|
log.Warn(ctx, "Query failed, failed to get shard delegator for query", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
|
|
// do query
|
|
results, err := sd.Query(queryCtx, req)
|
|
if err != nil {
|
|
log.Warn(ctx, "failed to query on delegator", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
|
|
// reduce result
|
|
tr.CtxElapse(ctx, fmt.Sprintf("start reduce query result, traceID = %s, vChannel = %s, segmentIDs = %v",
|
|
traceID,
|
|
channel,
|
|
req.GetSegmentIDs(),
|
|
))
|
|
|
|
if !node.manager.Collection.Ref(req.Req.GetCollectionID(), 1) {
|
|
err := merr.WrapErrCollectionNotFound(req.Req.GetCollectionID())
|
|
log.Warn(ctx, "Query failed, failed to get collection", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
collection := node.manager.Collection.Get(req.Req.GetCollectionID())
|
|
defer func() {
|
|
node.manager.Collection.Unref(req.GetReq().GetCollectionID(), 1)
|
|
}()
|
|
|
|
resp, err := segments.RunDelegatorQueryPipeline(ctx, req, collection.Schema(), results)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// aggregate cost
|
|
requestCosts := lo.FilterMap(results, func(result *internalpb.RetrieveResults, _ int) (*internalpb.CostAggregation, bool) {
|
|
if paramtable.Get().QueryNodeCfg.EnableWorkerSQCostMetrics.GetAsBool() {
|
|
return result.GetCostAggregation(), true
|
|
}
|
|
|
|
if result.GetBase().GetSourceID() == paramtable.GetNodeID() {
|
|
return result.GetCostAggregation(), true
|
|
}
|
|
|
|
return nil, false
|
|
})
|
|
resp.CostAggregation = segmentutil.MergeRequestCost(requestCosts)
|
|
if resp.CostAggregation == nil {
|
|
resp.CostAggregation = &internalpb.CostAggregation{}
|
|
}
|
|
relatedDataSize := lo.SumBy(results, func(t *internalpb.RetrieveResults) int64 {
|
|
cost := t.GetCostAggregation()
|
|
if cost == nil {
|
|
return 0
|
|
}
|
|
return cost.GetTotalRelatedDataSize()
|
|
})
|
|
resp.CostAggregation.TotalRelatedDataSize = relatedDataSize
|
|
tr.CtxElapse(ctx, fmt.Sprintf("do query with channel done , vChannel = %s, segmentIDs = %v",
|
|
channel,
|
|
req.GetSegmentIDs(),
|
|
))
|
|
|
|
latency := tr.ElapseSpan()
|
|
metrics.QueryNodeSQReqLatency.WithLabelValues(fmt.Sprint(node.GetNodeID()), queryLabel, metrics.Leader).Observe(float64(latency.Milliseconds()))
|
|
metrics.QueryNodeSQCount.WithLabelValues(fmt.Sprint(node.GetNodeID()), queryLabel, metrics.SuccessLabel, metrics.Leader, fmt.Sprint(req.GetReq().GetCollectionID())).Inc()
|
|
return resp, nil
|
|
}
|
|
|
|
func (node *QueryNode) queryChannelStream(ctx context.Context, req *querypb.QueryRequest, channel string, srv streamrpc.QueryStreamServer) error {
|
|
queryLabel := req.GetReq().GetQueryLabel()
|
|
if queryLabel == "" {
|
|
queryLabel = metrics.QueryLabel
|
|
}
|
|
ctx = contextutil.WithQueryLabel(ctx, queryLabel)
|
|
metrics.QueryNodeSQCount.WithLabelValues(fmt.Sprint(node.GetNodeID()), queryLabel, metrics.TotalLabel, metrics.Leader, fmt.Sprint(req.GetReq().GetCollectionID())).Inc()
|
|
msgID := req.Req.Base.GetMsgID()
|
|
log := mlog.With(
|
|
mlog.Int64("msgID", msgID),
|
|
mlog.FieldCollectionID(req.GetReq().GetCollectionID()),
|
|
mlog.String("channel", channel),
|
|
mlog.String("scope", req.GetScope().String()),
|
|
mlog.String("queryLabel", queryLabel),
|
|
)
|
|
|
|
var err error
|
|
defer func() {
|
|
if err != nil {
|
|
metrics.QueryNodeSQCount.WithLabelValues(fmt.Sprint(node.GetNodeID()), queryLabel, metrics.FailLabel, metrics.Leader, fmt.Sprint(req.GetReq().GetCollectionID())).Inc()
|
|
}
|
|
}()
|
|
|
|
log.Debug(ctx, "start do streaming query with channel",
|
|
mlog.Int64s("segmentIDs", req.GetSegmentIDs()),
|
|
)
|
|
|
|
// add cancel when error occurs
|
|
queryCtx, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
|
|
// From Proxy
|
|
tr := timerecord.NewTimeRecorder("queryDelegator")
|
|
// get delegator
|
|
sd, ok := node.delegators.Get(channel)
|
|
if !ok {
|
|
err := merr.WrapErrChannelNotFound(channel)
|
|
log.Warn(ctx, "Query failed, failed to get query shard delegator", mlog.Err(err))
|
|
return err
|
|
}
|
|
|
|
// do query
|
|
err = sd.QueryStream(queryCtx, req, srv)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
tr.CtxElapse(ctx, fmt.Sprintf("do query with channel done , vChannel = %s, segmentIDs = %v",
|
|
channel,
|
|
req.GetSegmentIDs(),
|
|
))
|
|
|
|
return nil
|
|
}
|
|
|
|
func (node *QueryNode) queryStreamSegments(ctx context.Context, req *querypb.QueryRequest, srv streamrpc.QueryStreamServer) error {
|
|
queryLabel := req.GetReq().GetQueryLabel()
|
|
if queryLabel == "" {
|
|
queryLabel = metrics.QueryLabel
|
|
}
|
|
ctx = contextutil.WithQueryLabel(ctx, queryLabel)
|
|
mlog.Debug(ctx, "received query stream request",
|
|
mlog.Int64s("outputFields", req.GetReq().GetOutputFieldsId()),
|
|
mlog.Int64s("segmentIDs", req.GetSegmentIDs()),
|
|
mlog.Uint64("guaranteeTimestamp", req.GetReq().GetGuaranteeTimestamp()),
|
|
mlog.Uint64("mvccTimestamp", req.GetReq().GetMvccTimestamp()),
|
|
)
|
|
|
|
if !node.manager.Collection.Ref(req.Req.GetCollectionID(), 1) {
|
|
err := merr.WrapErrCollectionNotFound(req.Req.GetCollectionID())
|
|
mlog.Warn(ctx, "Query stream segments failed, failed to get collection", mlog.Err(err))
|
|
return err
|
|
}
|
|
collection := node.manager.Collection.Get(req.Req.GetCollectionID())
|
|
defer func() {
|
|
node.manager.Collection.Unref(req.GetReq().GetCollectionID(), 1)
|
|
}()
|
|
|
|
// Send task to scheduler and wait until it finished.
|
|
|
|
task := tasks.NewQueryStreamTask(ctx, collection, node.manager, req, srv,
|
|
paramtable.Get().QueryNodeCfg.QueryStreamBatchSize.GetAsInt(),
|
|
paramtable.Get().QueryNodeCfg.QueryStreamMaxBatchSize.GetAsInt())
|
|
if err := node.scheduler.Add(task); err != nil {
|
|
mlog.Warn(ctx, "failed to add query task into scheduler", mlog.Err(err))
|
|
return err
|
|
}
|
|
|
|
err := task.Wait()
|
|
if err != nil {
|
|
mlog.Warn(ctx, "failed to execute task by node scheduler", mlog.Err(err))
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (node *QueryNode) searchChannel(ctx context.Context, req *querypb.SearchRequest, channel string) (*internalpb.SearchResults, error) {
|
|
log := mlog.With(
|
|
mlog.Int64("msgID", req.GetReq().GetBase().GetMsgID()),
|
|
mlog.FieldCollectionID(req.Req.GetCollectionID()),
|
|
mlog.String("channel", channel),
|
|
mlog.String("scope", req.GetScope().String()),
|
|
mlog.Int64("nq", req.GetReq().GetNq()),
|
|
)
|
|
|
|
if err := node.lifetime.Add(merr.IsHealthy); err != nil {
|
|
return nil, err
|
|
}
|
|
defer node.lifetime.Done()
|
|
|
|
nodeIDStr := paramtable.GetStringNodeID()
|
|
collIDStr := strconv.FormatInt(req.GetReq().GetCollectionID(), 10)
|
|
|
|
var err error
|
|
metrics.QueryNodeSQCount.WithLabelValues(nodeIDStr, metrics.SearchLabel, metrics.TotalLabel, metrics.Leader, collIDStr).Inc()
|
|
defer func() {
|
|
if err != nil {
|
|
metrics.QueryNodeSQCount.WithLabelValues(nodeIDStr, metrics.SearchLabel, metrics.FailLabel, metrics.Leader, collIDStr).Inc()
|
|
}
|
|
}()
|
|
|
|
log.Debug(ctx, "start to search channel",
|
|
mlog.Int64s("segmentIDs", req.GetSegmentIDs()),
|
|
)
|
|
|
|
// From Proxy
|
|
tr := timerecord.NewTimeRecorder("searchDelegator")
|
|
// get delegator
|
|
sd, ok := node.delegators.Get(channel)
|
|
if !ok {
|
|
err := merr.WrapErrChannelNotFound(channel)
|
|
log.Warn(ctx, "Query failed, failed to get shard delegator for search", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
// do search
|
|
results, err := sd.Search(ctx, req)
|
|
if err != nil {
|
|
log.Warn(ctx, "failed to search on delegator", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
|
|
tr.CtxElapse(ctx, "start reduce query result, ch="+channel)
|
|
|
|
resp, err := segments.ReduceSearchOnQueryNode(ctx, results,
|
|
reduce.NewReduceSearchResultInfo(req.GetReq().GetNq(),
|
|
req.GetReq().GetTopk()).WithMetricType(req.GetReq().GetMetricType()).
|
|
WithGroupSize(req.GetReq().GetGroupSize()).
|
|
WithGroupByFieldIdsFromProto(req.GetReq().GetGroupByFieldId(), req.GetReq().GetGroupByFieldIds()).
|
|
WithAdvance(req.GetReq().GetIsAdvanced()))
|
|
|
|
reduceLatency := tr.RecordSpan()
|
|
metrics.QueryNodeReduceLatency.
|
|
WithLabelValues(nodeIDStr, metrics.SearchLabel, metrics.ReduceShards, metrics.BatchReduce).
|
|
Observe(float64(reduceLatency.Microseconds()) / 1000.0)
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
tr.CtxElapse(ctx, "search with channel done, ch="+channel)
|
|
|
|
// update metric to prometheus
|
|
latency := tr.ElapseSpan()
|
|
metrics.QueryNodeSQReqLatency.WithLabelValues(nodeIDStr, metrics.SearchLabel, metrics.Leader).Observe(float64(latency.Milliseconds()))
|
|
metrics.QueryNodeSQCount.WithLabelValues(nodeIDStr, metrics.SearchLabel, metrics.SuccessLabel, metrics.Leader, collIDStr).Inc()
|
|
metrics.QueryNodeSearchNQ.WithLabelValues(nodeIDStr).Observe(float64(req.Req.GetNq()))
|
|
metrics.QueryNodeSearchTopK.WithLabelValues(nodeIDStr).Observe(float64(req.Req.GetTopk()))
|
|
return resp, nil
|
|
}
|
|
|
|
func (node *QueryNode) getChannelStatistics(ctx context.Context, req *querypb.GetStatisticsRequest, channel string) (*internalpb.GetStatisticsResponse, error) {
|
|
log := mlog.With(
|
|
mlog.FieldCollectionID(req.Req.GetCollectionID()),
|
|
mlog.String("channel", channel),
|
|
mlog.String("scope", req.GetScope().String()),
|
|
)
|
|
|
|
resp := &internalpb.GetStatisticsResponse{}
|
|
|
|
if req.GetFromShardLeader() {
|
|
var (
|
|
results []segments.SegmentStats
|
|
readSegments []segments.Segment
|
|
err error
|
|
)
|
|
|
|
switch req.GetScope() {
|
|
case querypb.DataScope_Historical:
|
|
results, readSegments, err = segments.StatisticsHistorical(ctx, node.manager, req.Req.GetCollectionID(), req.Req.GetPartitionIDs(), req.GetSegmentIDs())
|
|
case querypb.DataScope_Streaming:
|
|
results, readSegments, err = segments.StatisticStreaming(ctx, node.manager, req.Req.GetCollectionID(), req.Req.GetPartitionIDs(), req.GetSegmentIDs())
|
|
}
|
|
|
|
defer node.manager.Segment.Unpin(readSegments)
|
|
if err != nil {
|
|
log.Warn(ctx, "get segments statistics failed", mlog.Err(err))
|
|
return nil, err
|
|
}
|
|
return segmentStatsResponse(results), nil
|
|
}
|
|
|
|
sd, ok := node.delegators.Get(channel)
|
|
if !ok {
|
|
err := merr.WrapErrChannelNotFound(channel, "failed to get channel statistics")
|
|
log.Warn(ctx, "GetStatistics failed, failed to get query shard delegator", mlog.Err(err))
|
|
resp.Status = merr.Status(err)
|
|
return resp, nil
|
|
}
|
|
|
|
results, err := sd.GetStatistics(ctx, req)
|
|
if err != nil {
|
|
log.Warn(ctx, "failed to get statistics from delegator", mlog.Err(err))
|
|
resp.Status = merr.Status(err)
|
|
return resp, nil
|
|
}
|
|
resp, err = reduceStatisticResponse(results)
|
|
if err != nil {
|
|
log.Warn(ctx, "failed to reduce channel statistics", mlog.Err(err))
|
|
resp.Status = merr.Status(err)
|
|
return resp, nil
|
|
}
|
|
|
|
return resp, nil
|
|
}
|
|
|
|
func segmentStatsResponse(segStats []segments.SegmentStats) *internalpb.GetStatisticsResponse {
|
|
var totalRowNum int64
|
|
for _, stats := range segStats {
|
|
totalRowNum += stats.RowCount
|
|
}
|
|
|
|
resultMap := make(map[string]string)
|
|
resultMap["row_count"] = strconv.FormatInt(totalRowNum, 10)
|
|
|
|
ret := &internalpb.GetStatisticsResponse{
|
|
Status: merr.Success(),
|
|
Stats: funcutil.Map2KeyValuePair(resultMap),
|
|
}
|
|
return ret
|
|
}
|
|
|
|
func reduceStatisticResponse(results []*internalpb.GetStatisticsResponse) (*internalpb.GetStatisticsResponse, error) {
|
|
mergedResults := map[string]interface{}{
|
|
"row_count": int64(0),
|
|
}
|
|
fieldMethod := map[string]func(string) error{
|
|
"row_count": func(str string) error {
|
|
count, err := strconv.ParseInt(str, 10, 64)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
mergedResults["row_count"] = mergedResults["row_count"].(int64) + count
|
|
return nil
|
|
},
|
|
}
|
|
|
|
for _, partialResult := range results {
|
|
for _, pair := range partialResult.Stats {
|
|
fn, ok := fieldMethod[pair.Key]
|
|
if !ok {
|
|
// Stats keys are produced by query nodes; an unrecognized key is an
|
|
// internal protocol mismatch (e.g. mixed-version upgrade), not input.
|
|
return nil, merr.WrapErrServiceInternalMsg("unknown statistic field: %s", pair.Key)
|
|
}
|
|
if err := fn(pair.Value); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
}
|
|
|
|
stringMap := make(map[string]string)
|
|
for k, v := range mergedResults {
|
|
stringMap[k] = fmt.Sprint(v)
|
|
}
|
|
|
|
ret := &internalpb.GetStatisticsResponse{
|
|
Status: merr.Success(),
|
|
Stats: funcutil.Map2KeyValuePair(stringMap),
|
|
}
|
|
return ret, nil
|
|
}
|