package reader /* #cgo CFLAGS: -I../core/include #cgo LDFLAGS: -L../core/lib -lmilvus_dog_segment -Wl,-rpath=../core/lib #include "collection_c.h" #include "partition_c.h" #include "segment_c.h" */ import "C" import ( "github.com/czs007/suvlim/errors" "github.com/czs007/suvlim/pulsar/client-go/schema" "strconv" "unsafe" ) const SegmentLifetime = 20000 const ( SegmentOpened = 0 SegmentClosed = 1 ) type Segment struct { SegmentPtr C.CSegmentBase SegmentId uint64 SegmentCloseTime uint64 } func (s *Segment) GetStatus() int { /*C.IsOpened bool IsOpened(CSegmentBase c_segment); */ var isOpened = C.IsOpened(s.SegmentPtr) if isOpened { return SegmentOpened } else { return SegmentClosed } } func (s *Segment) GetSegmentID() uint64 { /*C.GetSegmentId unsigned long GetSegmentId(CSegmentBase c_segment); */ var segmentID = C.GetSegmentId(s.SegmentPtr) return uint64(segmentID) } func (s *Segment) SetSegmentID(segmentID uint64) { /*C.SetSegmentId void SetSegmentId(CSegmentBase c_segment, unsigned long segment_id); */ C.SetSegmentId(s.SegmentPtr, C.ulong(segmentID)) } func (s *Segment) GetMaxTimestamp() uint64 { /*C.GetTimeEnd unsigned long GetTimeEnd(CSegmentBase c_segment); */ var maxTimestamp = C.GetTimeEnd(s.SegmentPtr) return uint64(maxTimestamp) } func (s *Segment) SetMaxTimestamp(maxTimestamp uint64) { /*C.SetTimeEnd void SetTimeEnd(CSegmentBase c_segment, unsigned long time_end); */ C.SetTimeEnd(s.SegmentPtr, C.ulong(maxTimestamp)) } func (s *Segment) GetMinTimestamp() uint64 { /*C.GetTimeBegin unsigned long GetTimeBegin(CSegmentBase c_segment); */ var minTimestamp = C.GetTimeBegin(s.SegmentPtr) return uint64(minTimestamp) } func (s *Segment) SetMinTimestamp(minTimestamp uint64) { /*C.SetTimeBegin void SetTimeBegin(CSegmentBase c_segment, unsigned long time_begin); */ C.SetTimeBegin(s.SegmentPtr, C.ulong(minTimestamp)) } func (s *Segment) GetRowCount() int64 { /*C.GetRowCount long int GetRowCount(CSegmentBase c_segment); */ var rowCount = C.GetRowCount(s.SegmentPtr) return int64(rowCount) } func (s *Segment) GetDeletedCount() int64 { /*C.GetDeletedCount long int GetDeletedCount(CSegmentBase c_segment); */ var deletedCount = C.GetDeletedCount(s.SegmentPtr) return int64(deletedCount) } func (s *Segment) Close() error { /*C.Close int Close(CSegmentBase c_segment); */ var status = C.Close(s.SegmentPtr) if status != 0 { return errors.New("Close segment failed, error code = " + strconv.Itoa(int(status))) } return nil } //////////////////////////////////////////////////////////////////////////// func SegmentInsert(segment *Segment, entityIds *[]uint64, timestamps *[]uint64, dataChunk [][]*schema.FieldValue) (ResultEntityIds, error) { /*C.Insert int Insert(CSegmentBase c_segment, signed long int size, const unsigned long* primary_keys, const unsigned long* timestamps, void* raw_data, int sizeof_per_row, signed long int count); */ // TODO: remove hard code schema // auto schema_tmp = std::make_shared(); // schema_tmp->AddField("fakeVec", DataType::VECTOR_FLOAT, 16); // schema_tmp->AddField("age", DataType::INT32); // TODO: remove hard code & fake dataChunk const DIM = 4 const N = 3 var vec = [DIM]float32{1.1, 2.2, 3.3, 4.4} var rawData []int8 for i := 0; i <= N; i++ { for _, ele := range vec { rawData=append(rawData, int8(ele)) } rawData=append(rawData, int8(i)) } const sizeofPerRow = 4 + DIM * 4 var status = C.Insert(segment.SegmentPtr, C.long(N), (*C.ulong)(&(*entityIds)[0]), (*C.ulong)(&(*timestamps)[0]), unsafe.Pointer(&rawData[0]), C.int(sizeofPerRow), C.long(N)) if status != 0 { return nil, errors.New("Insert failed, error code = " + strconv.Itoa(int(status))) } return ResultEntityIds{}, nil } func SegmentDelete(segment *Segment, entityIds *[]uint64, timestamps *[]uint64) (ResultEntityIds, error) { /*C.Delete int Delete(CSegmentBase c_segment, long size, const unsigned long* primary_keys, const unsigned long* timestamps); */ size := len(*entityIds) var status = C.Delete(segment.SegmentPtr, C.long(size), (*C.ulong)(&(*entityIds)[0]), (*C.ulong)(&(*timestamps)[0])) if status != 0 { return nil, errors.New("Delete failed, error code = " + strconv.Itoa(int(status))) } return ResultEntityIds{}, nil } func SegmentSearch(segment *Segment, queryString string, timestamps *[]uint64, vectorRecord *[]schema.VectorRecord) (*[]SearchResult, error) { /*C.Search int Search(CSegmentBase c_segment, void* fake_query, unsigned long timestamp, long int* result_ids, float* result_distances); */ var results []SearchResult // TODO: get top-k's k from queryString const TopK = 1 for timestamp := range *timestamps { resultIds := make([]int64, TopK) resultDistances := make([]float32, TopK) var status = C.Search(segment.SegmentPtr, unsafe.Pointer(nil), C.ulong(timestamp), (*C.long)(&resultIds[0]), (*C.float)(&resultDistances[0])) if status != 0 { return nil, errors.New("Search failed, error code = " + strconv.Itoa(int(status))) } results = append(results, SearchResult{ResultIds: resultIds, ResultDistances: resultDistances}) } return &results, nil }