2021-04-19 13:47:10 +08:00
|
|
|
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
|
|
|
//
|
|
|
|
// Licensed 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.
|
|
|
|
|
2021-01-16 10:12:14 +08:00
|
|
|
package querynode
|
2020-11-26 16:01:31 +08:00
|
|
|
|
|
|
|
/*
|
|
|
|
#cgo CFLAGS: -I${SRCDIR}/../core/output/include
|
|
|
|
#cgo LDFLAGS: -L${SRCDIR}/../core/output/lib -lmilvus_segcore -Wl,-rpath=${SRCDIR}/../core/output/lib
|
|
|
|
|
|
|
|
#include "segcore/plan_c.h"
|
|
|
|
#include "segcore/reduce_c.h"
|
|
|
|
|
|
|
|
*/
|
|
|
|
import "C"
|
|
|
|
import (
|
2021-03-30 22:16:58 +08:00
|
|
|
"errors"
|
|
|
|
"fmt"
|
2020-12-03 19:00:11 +08:00
|
|
|
"strconv"
|
2021-03-30 22:16:58 +08:00
|
|
|
"sync"
|
2020-11-26 16:01:31 +08:00
|
|
|
"unsafe"
|
2021-03-05 10:15:27 +08:00
|
|
|
|
2021-04-22 14:45:57 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/log"
|
2020-11-26 16:01:31 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
type SearchResult struct {
|
2021-07-13 22:20:33 +08:00
|
|
|
cSearchResult C.CSearchResult
|
2020-11-26 16:01:31 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
type MarshaledHits struct {
|
|
|
|
cMarshaledHits C.CMarshaledHits
|
|
|
|
}
|
|
|
|
|
2020-12-03 19:00:11 +08:00
|
|
|
func reduceSearchResults(searchResults []*SearchResult, numSegments int64, inReduced []bool) error {
|
2021-07-13 22:20:33 +08:00
|
|
|
cSearchResults := make([]C.CSearchResult, 0)
|
2020-11-26 16:01:31 +08:00
|
|
|
for _, res := range searchResults {
|
2021-07-13 22:20:33 +08:00
|
|
|
cSearchResults = append(cSearchResults, res.cSearchResult)
|
2020-11-26 16:01:31 +08:00
|
|
|
}
|
2021-07-13 22:20:33 +08:00
|
|
|
cSearchResultPtr := (*C.CSearchResult)(&cSearchResults[0])
|
2020-11-26 16:01:31 +08:00
|
|
|
cNumSegments := C.long(numSegments)
|
2020-12-03 19:00:11 +08:00
|
|
|
cInReduced := (*C.bool)(&inReduced[0])
|
|
|
|
|
2021-07-13 22:20:33 +08:00
|
|
|
status := C.ReduceSearchResults(cSearchResultPtr, cNumSegments, cInReduced)
|
2020-12-03 19:00:11 +08:00
|
|
|
|
|
|
|
errorCode := status.error_code
|
|
|
|
|
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
|
|
|
return errors.New("reduceSearchResults failed, C runtime error detected, error code = " + strconv.Itoa(int(errorCode)) + ", error msg = " + errorMsg)
|
|
|
|
}
|
|
|
|
return nil
|
2020-11-26 16:01:31 +08:00
|
|
|
}
|
|
|
|
|
2021-07-13 22:20:33 +08:00
|
|
|
func fillTargetEntry(plan *SearchPlan, searchResults []*SearchResult, matchedSegments []*Segment, inReduced []bool) error {
|
2021-03-30 22:16:58 +08:00
|
|
|
wg := &sync.WaitGroup{}
|
|
|
|
fmt.Println(inReduced)
|
|
|
|
for i := range inReduced {
|
|
|
|
if inReduced[i] {
|
|
|
|
wg.Add(1)
|
|
|
|
go func(i int) {
|
|
|
|
err := matchedSegments[i].fillTargetEntry(plan, searchResults[i])
|
|
|
|
if err != nil {
|
|
|
|
log.Error(err.Error())
|
|
|
|
}
|
|
|
|
wg.Done()
|
|
|
|
}(i)
|
2020-12-03 19:00:11 +08:00
|
|
|
}
|
|
|
|
}
|
2021-03-30 22:16:58 +08:00
|
|
|
wg.Wait()
|
2020-12-03 19:00:11 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2021-07-13 22:20:33 +08:00
|
|
|
func reorganizeSearchResults(plan *SearchPlan, searchRequests []*searchRequest, searchResults []*SearchResult, numSegments int64, inReduced []bool) (*MarshaledHits, error) {
|
2020-11-26 16:01:31 +08:00
|
|
|
cPlaceholderGroups := make([]C.CPlaceholderGroup, 0)
|
2021-03-30 22:16:58 +08:00
|
|
|
for _, pg := range searchRequests {
|
2020-11-26 16:01:31 +08:00
|
|
|
cPlaceholderGroups = append(cPlaceholderGroups, (*pg).cPlaceholderGroup)
|
|
|
|
}
|
2020-12-03 19:00:11 +08:00
|
|
|
var cPlaceHolderGroupPtr = (*C.CPlaceholderGroup)(&cPlaceholderGroups[0])
|
2021-03-30 22:16:58 +08:00
|
|
|
var cNumGroup = (C.long)(len(searchRequests))
|
2020-12-03 19:00:11 +08:00
|
|
|
|
2021-07-13 22:20:33 +08:00
|
|
|
cSearchResults := make([]C.CSearchResult, 0)
|
2020-12-03 19:00:11 +08:00
|
|
|
for _, res := range searchResults {
|
2021-07-13 22:20:33 +08:00
|
|
|
cSearchResults = append(cSearchResults, res.cSearchResult)
|
2020-12-03 19:00:11 +08:00
|
|
|
}
|
2021-07-13 22:20:33 +08:00
|
|
|
cSearchResultPtr := (*C.CSearchResult)(&cSearchResults[0])
|
2020-12-03 19:00:11 +08:00
|
|
|
|
|
|
|
var cNumSegments = C.long(numSegments)
|
|
|
|
var cInReduced = (*C.bool)(&inReduced[0])
|
|
|
|
var cMarshaledHits C.CMarshaledHits
|
|
|
|
|
2021-07-13 22:20:33 +08:00
|
|
|
status := C.ReorganizeSearchResults(&cMarshaledHits, cPlaceHolderGroupPtr, cNumGroup, cSearchResultPtr, cInReduced, cNumSegments, plan.cSearchPlan)
|
2020-12-03 19:00:11 +08:00
|
|
|
errorCode := status.error_code
|
|
|
|
|
2021-03-30 22:16:58 +08:00
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
2021-07-13 22:20:33 +08:00
|
|
|
return nil, errors.New("reorganizeSearchResults failed, C runtime error detected, error code = " + strconv.Itoa(int(errorCode)) + ", error msg = " + errorMsg)
|
2021-03-30 22:16:58 +08:00
|
|
|
}
|
|
|
|
return &MarshaledHits{cMarshaledHits: cMarshaledHits}, nil
|
|
|
|
}
|
|
|
|
|
2021-07-13 22:20:33 +08:00
|
|
|
func reorganizeSingleSearchResult(plan *SearchPlan, placeholderGroups []*searchRequest, searchResult *SearchResult) (*MarshaledHits, error) {
|
2021-03-30 22:16:58 +08:00
|
|
|
cPlaceholderGroups := make([]C.CPlaceholderGroup, 0)
|
|
|
|
for _, pg := range placeholderGroups {
|
|
|
|
cPlaceholderGroups = append(cPlaceholderGroups, (*pg).cPlaceholderGroup)
|
|
|
|
}
|
|
|
|
var cPlaceHolderGroupPtr = (*C.CPlaceholderGroup)(&cPlaceholderGroups[0])
|
|
|
|
var cNumGroup = (C.long)(len(placeholderGroups))
|
|
|
|
|
2021-07-13 22:20:33 +08:00
|
|
|
cSearchResult := searchResult.cSearchResult
|
2021-03-30 22:16:58 +08:00
|
|
|
var cMarshaledHits C.CMarshaledHits
|
|
|
|
|
2021-07-13 22:20:33 +08:00
|
|
|
status := C.ReorganizeSingleSearchResult(&cMarshaledHits, cPlaceHolderGroupPtr, cNumGroup, cSearchResult, plan.cSearchPlan)
|
2021-03-30 22:16:58 +08:00
|
|
|
errorCode := status.error_code
|
|
|
|
|
2020-12-03 19:00:11 +08:00
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
2021-07-13 22:20:33 +08:00
|
|
|
return nil, errors.New("reorganizeSearchResults failed, C runtime error detected, error code = " + strconv.Itoa(int(errorCode)) + ", error msg = " + errorMsg)
|
2020-12-03 19:00:11 +08:00
|
|
|
}
|
|
|
|
return &MarshaledHits{cMarshaledHits: cMarshaledHits}, nil
|
2020-11-26 16:01:31 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (mh *MarshaledHits) getHitsBlobSize() int64 {
|
|
|
|
res := C.GetHitsBlobSize(mh.cMarshaledHits)
|
|
|
|
return int64(res)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (mh *MarshaledHits) getHitsBlob() ([]byte, error) {
|
|
|
|
byteSize := mh.getHitsBlobSize()
|
|
|
|
result := make([]byte, byteSize)
|
|
|
|
cResultPtr := unsafe.Pointer(&result[0])
|
|
|
|
C.GetHitsBlob(mh.cMarshaledHits, cResultPtr)
|
|
|
|
return result, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (mh *MarshaledHits) hitBlobSizeInGroup(groupOffset int64) ([]int64, error) {
|
|
|
|
cGroupOffset := (C.long)(groupOffset)
|
|
|
|
numQueries := C.GetNumQueriesPeerGroup(mh.cMarshaledHits, cGroupOffset)
|
|
|
|
result := make([]int64, int64(numQueries))
|
|
|
|
cResult := (*C.long)(&result[0])
|
|
|
|
C.GetHitSizePeerQueries(mh.cMarshaledHits, cGroupOffset, cResult)
|
|
|
|
return result, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func deleteMarshaledHits(hits *MarshaledHits) {
|
|
|
|
C.DeleteMarshaledHits(hits.cMarshaledHits)
|
|
|
|
}
|
|
|
|
|
|
|
|
func deleteSearchResults(results []*SearchResult) {
|
|
|
|
for _, result := range results {
|
2021-07-13 22:20:33 +08:00
|
|
|
C.DeleteSearchResult(result.cSearchResult)
|
2020-11-26 16:01:31 +08:00
|
|
|
}
|
|
|
|
}
|