2020-11-26 16:01:31 +08:00
|
|
|
package querynode
|
2020-08-25 15:45:19 +08:00
|
|
|
|
2020-08-29 17:42:41 +08:00
|
|
|
/*
|
|
|
|
|
2020-10-23 18:01:24 +08:00
|
|
|
#cgo CFLAGS: -I${SRCDIR}/../core/output/include
|
2020-08-29 17:42:41 +08:00
|
|
|
|
2020-10-31 15:11:47 +08:00
|
|
|
#cgo LDFLAGS: -L${SRCDIR}/../core/output/lib -lmilvus_segcore -Wl,-rpath=${SRCDIR}/../core/output/lib
|
2020-08-29 17:42:41 +08:00
|
|
|
|
2020-11-25 10:31:51 +08:00
|
|
|
#include "segcore/collection_c.h"
|
|
|
|
#include "segcore/plan_c.h"
|
|
|
|
#include "segcore/reduce_c.h"
|
2020-08-29 17:42:41 +08:00
|
|
|
*/
|
2020-08-25 15:45:19 +08:00
|
|
|
import "C"
|
2020-08-28 17:29:26 +08:00
|
|
|
import (
|
2020-11-04 17:58:43 +08:00
|
|
|
"strconv"
|
2020-11-05 10:52:50 +08:00
|
|
|
"unsafe"
|
2020-11-04 17:58:43 +08:00
|
|
|
|
2020-11-09 16:27:11 +08:00
|
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
|
2020-10-19 18:31:00 +08:00
|
|
|
"github.com/zilliztech/milvus-distributed/internal/errors"
|
2020-11-04 16:28:14 +08:00
|
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
|
2020-08-28 17:29:26 +08:00
|
|
|
)
|
|
|
|
|
2020-08-25 15:45:19 +08:00
|
|
|
type Segment struct {
|
2020-11-09 16:27:11 +08:00
|
|
|
segmentPtr C.CSegmentBase
|
|
|
|
segmentID UniqueID
|
2020-12-24 20:55:40 +08:00
|
|
|
partitionTag string // TODO: use partitionID
|
|
|
|
collectionID UniqueID
|
2020-11-09 16:27:11 +08:00
|
|
|
lastMemSize int64
|
|
|
|
lastRowCount int64
|
2020-11-05 10:52:50 +08:00
|
|
|
recentlyModified bool
|
2020-08-25 15:45:19 +08:00
|
|
|
}
|
|
|
|
|
2020-11-09 16:27:11 +08:00
|
|
|
func (s *Segment) ID() UniqueID {
|
|
|
|
return s.segmentID
|
|
|
|
}
|
|
|
|
|
|
|
|
//-------------------------------------------------------------------------------------- constructor and destructor
|
2020-12-24 20:55:40 +08:00
|
|
|
func newSegment(collection *Collection, segmentID int64, partitionTag string, collectionID UniqueID) *Segment {
|
2020-11-09 16:27:11 +08:00
|
|
|
/*
|
|
|
|
CSegmentBase
|
|
|
|
newSegment(CPartition partition, unsigned long segment_id);
|
|
|
|
*/
|
2020-11-09 18:07:15 +08:00
|
|
|
segmentPtr := C.NewSegment(collection.collectionPtr, C.ulong(segmentID))
|
2020-12-24 20:55:40 +08:00
|
|
|
var newSegment = &Segment{
|
|
|
|
segmentPtr: segmentPtr,
|
|
|
|
segmentID: segmentID,
|
|
|
|
partitionTag: partitionTag,
|
|
|
|
collectionID: collectionID,
|
|
|
|
}
|
2020-11-09 16:27:11 +08:00
|
|
|
|
|
|
|
return newSegment
|
|
|
|
}
|
|
|
|
|
|
|
|
func deleteSegment(segment *Segment) {
|
|
|
|
/*
|
|
|
|
void
|
|
|
|
deleteSegment(CSegmentBase segment);
|
|
|
|
*/
|
|
|
|
cPtr := segment.segmentPtr
|
|
|
|
C.DeleteSegment(cPtr)
|
|
|
|
}
|
|
|
|
|
|
|
|
//-------------------------------------------------------------------------------------- stats functions
|
2020-11-05 10:52:50 +08:00
|
|
|
func (s *Segment) getRowCount() int64 {
|
2020-09-21 18:16:06 +08:00
|
|
|
/*
|
2020-10-24 10:45:57 +08:00
|
|
|
long int
|
2020-11-05 10:52:50 +08:00
|
|
|
getRowCount(CSegmentBase c_segment);
|
2020-09-03 19:58:33 +08:00
|
|
|
*/
|
2020-11-09 16:27:11 +08:00
|
|
|
var rowCount = C.GetRowCount(s.segmentPtr)
|
2020-09-03 19:58:33 +08:00
|
|
|
return int64(rowCount)
|
2020-08-25 15:45:19 +08:00
|
|
|
}
|
|
|
|
|
2020-11-05 10:52:50 +08:00
|
|
|
func (s *Segment) getDeletedCount() int64 {
|
2020-09-21 18:16:06 +08:00
|
|
|
/*
|
2020-10-24 10:45:57 +08:00
|
|
|
long int
|
2020-11-05 10:52:50 +08:00
|
|
|
getDeletedCount(CSegmentBase c_segment);
|
2020-09-03 19:58:33 +08:00
|
|
|
*/
|
2020-11-09 16:27:11 +08:00
|
|
|
var deletedCount = C.GetDeletedCount(s.segmentPtr)
|
2020-09-03 19:58:33 +08:00
|
|
|
return int64(deletedCount)
|
|
|
|
}
|
|
|
|
|
2020-11-05 10:52:50 +08:00
|
|
|
func (s *Segment) getMemSize() int64 {
|
2020-09-21 18:16:06 +08:00
|
|
|
/*
|
2020-10-24 10:45:57 +08:00
|
|
|
long int
|
|
|
|
GetMemoryUsageInBytes(CSegmentBase c_segment);
|
2020-09-21 15:10:54 +08:00
|
|
|
*/
|
2020-11-09 16:27:11 +08:00
|
|
|
var memoryUsageInBytes = C.GetMemoryUsageInBytes(s.segmentPtr)
|
2020-09-21 15:10:54 +08:00
|
|
|
|
2020-10-29 19:55:57 +08:00
|
|
|
return int64(memoryUsageInBytes)
|
2020-09-16 15:21:10 +08:00
|
|
|
}
|
|
|
|
|
2020-11-12 11:18:23 +08:00
|
|
|
//-------------------------------------------------------------------------------------- preDm functions
|
2020-11-05 10:52:50 +08:00
|
|
|
func (s *Segment) segmentPreInsert(numOfRecords int) int64 {
|
2020-09-21 18:16:06 +08:00
|
|
|
/*
|
2020-10-24 10:45:57 +08:00
|
|
|
long int
|
|
|
|
PreInsert(CSegmentBase c_segment, long int size);
|
2020-09-09 15:24:07 +08:00
|
|
|
*/
|
2020-11-09 16:27:11 +08:00
|
|
|
var offset = C.PreInsert(s.segmentPtr, C.long(int64(numOfRecords)))
|
2020-09-09 15:24:07 +08:00
|
|
|
|
2020-09-12 16:57:37 +08:00
|
|
|
return int64(offset)
|
2020-09-09 15:24:07 +08:00
|
|
|
}
|
|
|
|
|
2020-11-05 10:52:50 +08:00
|
|
|
func (s *Segment) segmentPreDelete(numOfRecords int) int64 {
|
2020-09-21 18:16:06 +08:00
|
|
|
/*
|
2020-10-24 10:45:57 +08:00
|
|
|
long int
|
|
|
|
PreDelete(CSegmentBase c_segment, long int size);
|
2020-09-09 15:24:07 +08:00
|
|
|
*/
|
2020-11-09 16:27:11 +08:00
|
|
|
var offset = C.PreDelete(s.segmentPtr, C.long(int64(numOfRecords)))
|
2020-09-09 15:24:07 +08:00
|
|
|
|
2020-09-12 16:57:37 +08:00
|
|
|
return int64(offset)
|
2020-09-09 15:24:07 +08:00
|
|
|
}
|
|
|
|
|
2020-11-09 16:27:11 +08:00
|
|
|
//-------------------------------------------------------------------------------------- dm & search functions
|
2020-11-05 10:52:50 +08:00
|
|
|
func (s *Segment) segmentInsert(offset int64, entityIDs *[]UniqueID, timestamps *[]Timestamp, records *[]*commonpb.Blob) error {
|
2020-09-21 18:16:06 +08:00
|
|
|
/*
|
2020-11-26 16:01:31 +08:00
|
|
|
CStatus
|
2020-10-24 10:45:57 +08:00
|
|
|
Insert(CSegmentBase c_segment,
|
|
|
|
long int reserved_offset,
|
|
|
|
signed long int size,
|
|
|
|
const long* primary_keys,
|
|
|
|
const unsigned long* timestamps,
|
|
|
|
void* raw_data,
|
|
|
|
int sizeof_per_row,
|
|
|
|
signed long int count);
|
2020-09-02 16:23:50 +08:00
|
|
|
*/
|
2020-09-08 10:39:09 +08:00
|
|
|
// Blobs to one big blob
|
2020-11-05 10:52:50 +08:00
|
|
|
var numOfRow = len(*entityIDs)
|
|
|
|
var sizeofPerRow = len((*records)[0].Value)
|
|
|
|
|
|
|
|
assert.Equal(nil, numOfRow, len(*records))
|
|
|
|
|
|
|
|
var rawData = make([]byte, numOfRow*sizeofPerRow)
|
|
|
|
var copyOffset = 0
|
|
|
|
for i := 0; i < len(*records); i++ {
|
|
|
|
copy(rawData[copyOffset:], (*records)[i].Value)
|
|
|
|
copyOffset += sizeofPerRow
|
|
|
|
}
|
|
|
|
|
|
|
|
var cOffset = C.long(offset)
|
|
|
|
var cNumOfRows = C.long(numOfRow)
|
|
|
|
var cEntityIdsPtr = (*C.long)(&(*entityIDs)[0])
|
|
|
|
var cTimestampsPtr = (*C.ulong)(&(*timestamps)[0])
|
|
|
|
var cSizeofPerRow = C.int(sizeofPerRow)
|
|
|
|
var cRawDataVoidPtr = unsafe.Pointer(&rawData[0])
|
|
|
|
|
2020-11-09 16:27:11 +08:00
|
|
|
var status = C.Insert(s.segmentPtr,
|
2020-11-05 10:52:50 +08:00
|
|
|
cOffset,
|
|
|
|
cNumOfRows,
|
|
|
|
cEntityIdsPtr,
|
|
|
|
cTimestampsPtr,
|
|
|
|
cRawDataVoidPtr,
|
|
|
|
cSizeofPerRow,
|
|
|
|
cNumOfRows)
|
|
|
|
|
2020-11-26 16:01:31 +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("Insert failed, C runtime error detected, error code = " + strconv.Itoa(int(errorCode)) + ", error msg = " + errorMsg)
|
2020-11-05 10:52:50 +08:00
|
|
|
}
|
2020-09-02 17:18:49 +08:00
|
|
|
|
2021-01-11 15:17:06 +08:00
|
|
|
s.recentlyModified = true
|
2020-09-07 17:01:46 +08:00
|
|
|
return nil
|
2020-08-28 17:29:26 +08:00
|
|
|
}
|
|
|
|
|
2020-11-05 10:52:50 +08:00
|
|
|
func (s *Segment) segmentDelete(offset int64, entityIDs *[]UniqueID, timestamps *[]Timestamp) error {
|
2020-09-21 18:16:06 +08:00
|
|
|
/*
|
2020-11-26 16:01:31 +08:00
|
|
|
CStatus
|
2020-10-24 10:45:57 +08:00
|
|
|
Delete(CSegmentBase c_segment,
|
|
|
|
long int reserved_offset,
|
|
|
|
long size,
|
|
|
|
const long* primary_keys,
|
|
|
|
const unsigned long* timestamps);
|
2020-09-02 16:23:50 +08:00
|
|
|
*/
|
2020-09-09 15:24:07 +08:00
|
|
|
var cOffset = C.long(offset)
|
|
|
|
var cSize = C.long(len(*entityIDs))
|
2020-09-12 16:57:37 +08:00
|
|
|
var cEntityIdsPtr = (*C.long)(&(*entityIDs)[0])
|
2020-09-09 15:24:07 +08:00
|
|
|
var cTimestampsPtr = (*C.ulong)(&(*timestamps)[0])
|
2020-09-02 16:23:50 +08:00
|
|
|
|
2020-11-09 16:27:11 +08:00
|
|
|
var status = C.Delete(s.segmentPtr, cOffset, cSize, cEntityIdsPtr, cTimestampsPtr)
|
2020-09-02 17:18:49 +08:00
|
|
|
|
2020-11-26 16:01:31 +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("Delete failed, C runtime error detected, error code = " + strconv.Itoa(int(errorCode)) + ", error msg = " + errorMsg)
|
2020-09-02 17:18:49 +08:00
|
|
|
}
|
2020-09-02 16:23:50 +08:00
|
|
|
|
2020-09-07 17:01:46 +08:00
|
|
|
return nil
|
2020-08-28 17:29:26 +08:00
|
|
|
}
|
|
|
|
|
2020-11-19 14:13:39 +08:00
|
|
|
func (s *Segment) segmentSearch(plan *Plan,
|
|
|
|
placeHolderGroups []*PlaceholderGroup,
|
2020-11-26 16:01:31 +08:00
|
|
|
timestamp []Timestamp) (*SearchResult, error) {
|
2020-11-09 16:27:11 +08:00
|
|
|
/*
|
2020-11-26 16:01:31 +08:00
|
|
|
CStatus
|
|
|
|
Search(void* plan,
|
2020-11-19 14:13:39 +08:00
|
|
|
void* placeholder_groups,
|
|
|
|
uint64_t* timestamps,
|
|
|
|
int num_groups,
|
|
|
|
long int* result_ids,
|
|
|
|
float* result_distances);
|
2020-11-09 16:27:11 +08:00
|
|
|
*/
|
|
|
|
|
2020-11-17 10:07:42 +08:00
|
|
|
cPlaceholderGroups := make([]C.CPlaceholderGroup, 0)
|
|
|
|
for _, pg := range placeHolderGroups {
|
|
|
|
cPlaceholderGroups = append(cPlaceholderGroups, (*pg).cPlaceholderGroup)
|
2020-09-18 01:53:18 +08:00
|
|
|
}
|
2020-09-02 16:23:50 +08:00
|
|
|
|
2020-11-26 16:01:31 +08:00
|
|
|
var searchResult SearchResult
|
2020-11-17 10:07:42 +08:00
|
|
|
var cTimestamp = (*C.ulong)(×tamp[0])
|
|
|
|
var cPlaceHolder = (*C.CPlaceholderGroup)(&cPlaceholderGroups[0])
|
|
|
|
var cNumGroups = C.int(len(placeHolderGroups))
|
2020-12-03 19:00:11 +08:00
|
|
|
var cQueryResult = (*C.CQueryResult)(&searchResult.cQueryResult)
|
2020-09-02 16:23:50 +08:00
|
|
|
|
2020-11-26 16:01:31 +08:00
|
|
|
var status = C.Search(s.segmentPtr, plan.cPlan, cPlaceHolder, cTimestamp, cNumGroups, cQueryResult)
|
|
|
|
errorCode := status.error_code
|
2020-11-26 14:29:01 +08:00
|
|
|
|
2020-11-26 16:01:31 +08:00
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
|
|
|
return nil, errors.New("Search failed, C runtime error detected, error code = " + strconv.Itoa(int(errorCode)) + ", error msg = " + errorMsg)
|
2020-11-26 15:18:36 +08:00
|
|
|
}
|
2020-11-26 16:01:31 +08:00
|
|
|
|
|
|
|
return &searchResult, nil
|
2020-08-28 17:29:26 +08:00
|
|
|
}
|
2020-12-03 19:00:11 +08:00
|
|
|
|
|
|
|
func (s *Segment) fillTargetEntry(plan *Plan,
|
|
|
|
result *SearchResult) error {
|
|
|
|
|
|
|
|
var status = C.FillTargetEntry(s.segmentPtr, plan.cPlan, result.cQueryResult)
|
|
|
|
errorCode := status.error_code
|
|
|
|
|
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
|
|
|
return errors.New("FillTargetEntry failed, C runtime error detected, error code = " + strconv.Itoa(int(errorCode)) + ", error msg = " + errorMsg)
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
2020-12-24 20:55:40 +08:00
|
|
|
|
|
|
|
// segment, err := loadIndexService.replica.getSegmentByID(segmentID)
|
|
|
|
func (s *Segment) updateSegmentIndex(loadIndexInfo *LoadIndexInfo) error {
|
|
|
|
status := C.UpdateSegmentIndex(s.segmentPtr, loadIndexInfo.cLoadIndexInfo)
|
|
|
|
errorCode := status.error_code
|
|
|
|
|
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
|
|
|
return errors.New("updateSegmentIndex failed, C runtime error detected, error code = " + strconv.Itoa(int(errorCode)) + ", error msg = " + errorMsg)
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|