milvus2.4多向量搜索源码分析

api入口

HybridSearch是多向量搜索的API。

func (node *Proxy) HybridSearch(ctx context.Context, request *milvuspb.HybridSearchRequest) (*milvuspb.SearchResults, error) {
	var err error
	rsp := &milvuspb.SearchResults{
		Status: merr.Success(),
	}
	err2 := retry.Handle(ctx, func() (bool, error) {
		rsp, err = node.hybridSearch(ctx, request)
		if errors.Is(merr.Error(rsp.GetStatus()), merr.ErrInconsistentRequery) {
			return true, merr.Error(rsp.GetStatus())
		}
		return false, nil
	})
	if err2 != nil {
		rsp.Status = merr.Status(err2)
	}
	return rsp, err
}

func (node *Proxy) hybridSearch(ctx context.Context, request *milvuspb.HybridSearchRequest) (*milvuspb.SearchResults, error) {
	......
    // 转换为milvuspb.SearchRequest
	newSearchReq := convertHybridSearchToSearch(request)
	qt := &searchTask{
		ctx:       ctx,
		Condition: NewTaskCondition(ctx),
		SearchRequest: &internalpb.SearchRequest{
			Base: commonpbutil.NewMsgBase(
				commonpbutil.WithMsgType(commonpb.MsgType_Search),
				commonpbutil.WithSourceID(paramtable.GetNodeID()),
			),
			ReqID: paramtable.GetNodeID(),
		},
		request:             newSearchReq,
		tr:                  timerecord.NewTimeRecorder(method),
		qc:                  node.queryCoord,
		node:                node,
		lb:                  node.lbPolicy,
		mustUsePartitionKey: Params.ProxyCfg.MustUsePartitionKey.GetAsBool(),
	}

	guaranteeTs := request.GuaranteeTimestamp

	log := log.Ctx(ctx).With(
		zap.String("role", typeutil.ProxyRole),
		zap.String("db", request.DbName),
		zap.String("collection", request.CollectionName),
		zap.Any("partitions", request.PartitionNames),
		zap.Any("OutputFields", request.OutputFields),
		zap.Uint64("guarantee_timestamp", guaranteeTs),
	)

	defer func() {
		span := tr.ElapseSpan()
		if span >= paramtable.Get().ProxyCfg.SlowQuerySpanInSeconds.GetAsDuration(time.Second) {
			log.Info(rpcSlow(method), zap.Duration("duration", span))
			metrics.ProxySlowQueryCount.WithLabelValues(
				strconv.FormatInt(paramtable.GetNodeID(), 10),
				metrics.HybridSearchLabel,
			).Inc()
		}
	}()

	log.Debug(rpcReceived(method))

	if err := node.sched.dqQueue.Enqueue(qt); err != nil {
		......
	}
	......
}

从代码中可以看出HybridSearch最终调用的和Search() API是同一个task。

milvuspb.HybridSearchRequest有一个[]*SearchRequest变量，这个存储了多个查询结构体，如果是普通的Search()，传参直接就是SearchRequest结构体，如果是HybridSearch()，就是多个查询结构体，下一步做转换。

convertHybridSearchToSearch

进入convertHybridSearchToSearch()看看是如何转换的。

func convertHybridSearchToSearch(req *milvuspb.HybridSearchRequest) *milvuspb.SearchRequest {
	ret := &milvuspb.SearchRequest{
		Base:                  req.GetBase(),
		DbName:                req.GetDbName(),
		CollectionName:        req.GetCollectionName(),
		PartitionNames:        req.GetPartitionNames(),
		OutputFields:          req.GetOutputFields(),
		SearchParams:          req.GetRankParams(),
		TravelTimestamp:       req.GetTravelTimestamp(),
		GuaranteeTimestamp:    req.GetGuaranteeTimestamp(),
		Nq:                    0,
		NotReturnAllMeta:      req.GetNotReturnAllMeta(),
		ConsistencyLevel:      req.GetConsistencyLevel(),
		UseDefaultConsistency: req.GetUseDefaultConsistency(),
		SearchByPrimaryKeys:   false,
		SubReqs:               nil,
	}

	for _, sub := range req.GetRequests() {
		subReq := &milvuspb.SubSearchRequest{
			Dsl:              sub.GetDsl(),
			PlaceholderGroup: sub.GetPlaceholderGroup(),
			DslType:          sub.GetDslType(),
			SearchParams:     sub.GetSearchParams(),
			Nq:               sub.GetNq(),
		}
		ret.SubReqs = append(ret.SubReqs, subReq)
	}
	return ret
}

milvuspb.SearchRequest结构体增加了一个SubReqs变量，类型是[]*SubSearchRequest。

type SubSearchRequest struct {
    Dsl                  string
    PlaceholderGroup     []byte
    DslType              commonpb.DslType
    SearchParams         []*commonpb.KeyValuePair
    Nq                   int64
    XXX_NoUnkeyedLiteral struct{}
    XXX_unrecognized     []byte
    XXX_sizecache        int32
}

searchTask

PreExecute()

t.SearchRequest.IsAdvanced = len(t.request.GetSubReqs()) > 0

Search()和HybridSearch()最终都是走的searchTask，如果是HybridSearch，IsAdvanced会置为true，如果是Search，IsAdvanced会置为false。

func (t *searchTask) PreExecute(ctx context.Context) error {
	......
	t.SearchRequest.IsAdvanced = len(t.request.GetSubReqs()) > 0
	......

	if t.SearchRequest.GetIsAdvanced() {
		if len(t.request.GetSubReqs()) > defaultMaxSearchRequest {
			return errors.New(fmt.Sprintf("maximum of ann search requests is %d", defaultMaxSearchRequest))
		}
	}
	if t.SearchRequest.GetIsAdvanced() {
		t.rankParams, err = parseRankParams(t.request.GetSearchParams())
		if err != nil {
			return err
		}
	}
	......

	if t.SearchRequest.GetIsAdvanced() {
		t.requery = len(t.request.OutputFields) > 0
		err = t.initAdvancedSearchRequest(ctx)
	} else {
		t.requery = len(vectorOutputFields) > 0
		err = t.initSearchRequest(ctx)
	}
	......
}

在initAdvancedSearchRequest填充SubReqs

Search

最终会转换为多次search。

// Search preforms search operation on shard.
func (sd *shardDelegator) Search(ctx context.Context, req *querypb.SearchRequest) ([]*internalpb.SearchResults, error) {
	......
	if req.GetReq().GetIsAdvanced() {
		futures := make([]*conc.Future[*internalpb.SearchResults], len(req.GetReq().GetSubReqs()))
        // 多次调用search
		for index, subReq := range req.GetReq().GetSubReqs() {
			newRequest := &internalpb.SearchRequest{
				Base:               req.GetReq().GetBase(),
				ReqID:              req.GetReq().GetReqID(),
				DbID:               req.GetReq().GetDbID(),
				CollectionID:       req.GetReq().GetCollectionID(),
				PartitionIDs:       subReq.GetPartitionIDs(),
				Dsl:                subReq.GetDsl(),
				PlaceholderGroup:   subReq.GetPlaceholderGroup(),
				DslType:            subReq.GetDslType(),
				SerializedExprPlan: subReq.GetSerializedExprPlan(),
				OutputFieldsId:     req.GetReq().GetOutputFieldsId(),
				MvccTimestamp:      req.GetReq().GetMvccTimestamp(),
				GuaranteeTimestamp: req.GetReq().GetGuaranteeTimestamp(),
				TimeoutTimestamp:   req.GetReq().GetTimeoutTimestamp(),
				Nq:                 subReq.GetNq(),
				Topk:               subReq.GetTopk(),
				MetricType:         subReq.GetMetricType(),
				IgnoreGrowing:      req.GetReq().GetIgnoreGrowing(),
				Username:           req.GetReq().GetUsername(),
				IsAdvanced:         false,
			}
			future := conc.Go(func() (*internalpb.SearchResults, error) {
				searchReq := &querypb.SearchRequest{
					Req:             newRequest,
					DmlChannels:     req.GetDmlChannels(),
					TotalChannelNum: req.GetTotalChannelNum(),
				}
				searchReq.Req.GuaranteeTimestamp = req.GetReq().GetGuaranteeTimestamp()
				searchReq.Req.TimeoutTimestamp = req.GetReq().GetTimeoutTimestamp()
				if searchReq.GetReq().GetMvccTimestamp() == 0 {
					searchReq.GetReq().MvccTimestamp = tSafe
				}
                // 执行搜索
				results, err := sd.search(ctx, searchReq, sealed, growing)
				if err != nil {
					return nil, err
				}

				return segments.ReduceSearchResults(ctx,
					results,
					searchReq.Req.GetNq(),
					searchReq.Req.GetTopk(),
					searchReq.Req.GetMetricType())
			})
			futures[index] = future
		}
        // 等待所有任务执行完成
		err = conc.AwaitAll(futures...)
		if err != nil {
			return nil, err
		}
		results := make([]*internalpb.SearchResults, len(futures))
		for i, future := range futures {
			result := future.Value()
			if result.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success {
				log.Debug("delegator hybrid search failed",
					zap.String("reason", result.GetStatus().GetReason()))
				return nil, merr.Error(result.GetStatus())
			}
			results[i] = result
		}
		var ret *internalpb.SearchResults
		ret, err = segments.MergeToAdvancedResults(ctx, results)
		if err != nil {
			return nil, err
		}
        // 走这里
		return []*internalpb.SearchResults{ret}, nil
	}
	return sd.search(ctx, req, sealed, growing)
}

总结

HybridSearch会转换为多个Search搜索。