2021-04-19 13:50:12 +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-15 14:38:36 +08:00
|
|
|
package indexnode
|
2020-12-18 15:44:27 +08:00
|
|
|
|
|
|
|
/*
|
|
|
|
|
|
|
|
#cgo CFLAGS: -I${SRCDIR}/../core/output/include
|
|
|
|
|
|
|
|
#cgo LDFLAGS: -L${SRCDIR}/../core/output/lib -lmilvus_indexbuilder -Wl,-rpath=${SRCDIR}/../core/output/lib
|
|
|
|
|
2020-12-21 16:36:07 +08:00
|
|
|
#include <stdlib.h> // free
|
2020-12-18 15:44:27 +08:00
|
|
|
#include "segcore/collection_c.h"
|
|
|
|
#include "indexbuilder/index_c.h"
|
|
|
|
|
|
|
|
*/
|
|
|
|
import "C"
|
2021-01-26 09:38:40 +08:00
|
|
|
|
2020-12-18 15:44:27 +08:00
|
|
|
import (
|
2021-01-12 18:03:24 +08:00
|
|
|
"fmt"
|
2020-12-21 16:36:07 +08:00
|
|
|
"unsafe"
|
2020-12-18 15:44:27 +08:00
|
|
|
|
2021-03-12 19:47:37 +08:00
|
|
|
"go.uber.org/zap"
|
|
|
|
|
2020-12-21 16:36:07 +08:00
|
|
|
"github.com/golang/protobuf/proto"
|
2021-04-22 14:45:57 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/log"
|
|
|
|
"github.com/milvus-io/milvus/internal/proto/commonpb"
|
|
|
|
"github.com/milvus-io/milvus/internal/proto/indexcgopb"
|
|
|
|
"github.com/milvus-io/milvus/internal/storage"
|
2020-12-18 15:44:27 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
// TODO: use storage.Blob instead later
|
2020-12-25 11:10:31 +08:00
|
|
|
type Blob = storage.Blob
|
2020-12-18 15:44:27 +08:00
|
|
|
|
|
|
|
type Index interface {
|
|
|
|
Serialize() ([]*Blob, error)
|
|
|
|
Load([]*Blob) error
|
2020-12-28 10:00:02 +08:00
|
|
|
BuildFloatVecIndexWithoutIds(vectors []float32) error
|
|
|
|
BuildBinaryVecIndexWithoutIds(vectors []byte) error
|
2020-12-18 15:44:27 +08:00
|
|
|
Delete() error
|
|
|
|
}
|
|
|
|
|
|
|
|
type CIndex struct {
|
|
|
|
indexPtr C.CIndex
|
|
|
|
}
|
|
|
|
|
|
|
|
func (index *CIndex) Serialize() ([]*Blob, error) {
|
2021-02-22 18:33:40 +08:00
|
|
|
var cBinary C.CBinary
|
|
|
|
|
|
|
|
status := C.SerializeToSlicedBuffer(index.indexPtr, &cBinary)
|
2021-02-23 14:35:00 +08:00
|
|
|
defer C.DeleteCBinary(cBinary)
|
2020-12-29 17:53:27 +08:00
|
|
|
errorCode := status.error_code
|
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
2021-03-05 10:15:27 +08:00
|
|
|
return nil, fmt.Errorf("SerializeToSlicedBuffer failed, C runtime error detected, error code = %d, err msg = %s", errorCode, errorMsg)
|
2020-12-29 17:53:27 +08:00
|
|
|
}
|
|
|
|
|
2021-02-22 18:33:40 +08:00
|
|
|
binarySize := C.GetCBinarySize(cBinary)
|
|
|
|
binaryData := make([]byte, binarySize)
|
|
|
|
C.GetCBinaryData(cBinary, unsafe.Pointer(&binaryData[0]))
|
2020-12-21 16:36:07 +08:00
|
|
|
|
|
|
|
var blobs indexcgopb.BinarySet
|
2021-02-22 18:33:40 +08:00
|
|
|
err := proto.Unmarshal(binaryData, &blobs)
|
2020-12-18 15:44:27 +08:00
|
|
|
if err != nil {
|
2020-12-21 16:36:07 +08:00
|
|
|
return nil, err
|
2020-12-18 15:44:27 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
ret := make([]*Blob, 0)
|
2020-12-21 16:36:07 +08:00
|
|
|
for _, data := range blobs.Datas {
|
|
|
|
ret = append(ret, &Blob{Key: data.Key, Value: data.Value})
|
2020-12-18 15:44:27 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
return ret, nil
|
|
|
|
}
|
|
|
|
|
2020-12-21 16:36:07 +08:00
|
|
|
func (index *CIndex) Load(blobs []*Blob) error {
|
|
|
|
binarySet := &indexcgopb.BinarySet{Datas: make([]*indexcgopb.Binary, 0)}
|
|
|
|
for _, blob := range blobs {
|
|
|
|
binarySet.Datas = append(binarySet.Datas, &indexcgopb.Binary{Key: blob.Key, Value: blob.Value})
|
|
|
|
}
|
|
|
|
|
2021-02-24 17:24:51 +08:00
|
|
|
datas, err2 := proto.Marshal(binarySet)
|
|
|
|
if err2 != nil {
|
|
|
|
return err2
|
2020-12-21 16:36:07 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
/*
|
2020-12-30 18:16:34 +08:00
|
|
|
CStatus
|
2020-12-25 11:10:31 +08:00
|
|
|
LoadFromSlicedBuffer(CIndex index, const char* serialized_sliced_blob_buffer, int32_t size);
|
2020-12-21 16:36:07 +08:00
|
|
|
*/
|
2020-12-30 18:16:34 +08:00
|
|
|
status := C.LoadFromSlicedBuffer(index.indexPtr, (*C.char)(unsafe.Pointer(&datas[0])), (C.int32_t)(len(datas)))
|
|
|
|
errorCode := status.error_code
|
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
2021-03-05 10:15:27 +08:00
|
|
|
return fmt.Errorf("BuildFloatVecIndexWithoutIds failed, C runtime error detected, error code = %d, err msg = %s", errorCode, errorMsg)
|
2020-12-30 18:16:34 +08:00
|
|
|
}
|
2020-12-18 15:44:27 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2020-12-28 10:00:02 +08:00
|
|
|
func (index *CIndex) BuildFloatVecIndexWithoutIds(vectors []float32) error {
|
2020-12-21 16:36:07 +08:00
|
|
|
/*
|
2020-12-29 17:53:27 +08:00
|
|
|
CStatus
|
2020-12-28 10:00:02 +08:00
|
|
|
BuildFloatVecIndexWithoutIds(CIndex index, int64_t float_value_num, const float* vectors);
|
2020-12-21 16:36:07 +08:00
|
|
|
*/
|
2021-03-12 19:47:37 +08:00
|
|
|
log.Debug("before BuildFloatVecIndexWithoutIds")
|
2020-12-29 17:53:27 +08:00
|
|
|
status := C.BuildFloatVecIndexWithoutIds(index.indexPtr, (C.int64_t)(len(vectors)), (*C.float)(&vectors[0]))
|
|
|
|
errorCode := status.error_code
|
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
2021-03-12 19:47:37 +08:00
|
|
|
log.Debug("indexnode", zap.String("BuildFloatVecIndexWithoutIds error msg: ", errorMsg))
|
2020-12-29 17:53:27 +08:00
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
2021-03-05 10:15:27 +08:00
|
|
|
return fmt.Errorf("BuildFloatVecIndexWithoutIds failed, C runtime error detected, error code = %d, err msg = %s", errorCode, errorMsg)
|
2020-12-29 17:53:27 +08:00
|
|
|
}
|
2020-12-28 10:00:02 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (index *CIndex) BuildBinaryVecIndexWithoutIds(vectors []byte) error {
|
|
|
|
/*
|
2020-12-29 17:53:27 +08:00
|
|
|
CStatus
|
2020-12-28 10:00:02 +08:00
|
|
|
BuildBinaryVecIndexWithoutIds(CIndex index, int64_t data_size, const uint8_t* vectors);
|
|
|
|
*/
|
2020-12-29 17:53:27 +08:00
|
|
|
status := C.BuildBinaryVecIndexWithoutIds(index.indexPtr, (C.int64_t)(len(vectors)), (*C.uint8_t)(&vectors[0]))
|
|
|
|
errorCode := status.error_code
|
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
2021-03-05 10:15:27 +08:00
|
|
|
return fmt.Errorf("BuildBinaryVecIndexWithoutIds failed, C runtime error detected, error code = %d, err msg = %s", errorCode, errorMsg)
|
2020-12-29 17:53:27 +08:00
|
|
|
}
|
2020-12-18 15:44:27 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (index *CIndex) Delete() error {
|
2020-12-21 16:36:07 +08:00
|
|
|
/*
|
|
|
|
void
|
|
|
|
DeleteIndex(CIndex index);
|
|
|
|
*/
|
2020-12-18 15:44:27 +08:00
|
|
|
C.DeleteIndex(index.indexPtr)
|
2021-01-04 12:03:29 +08:00
|
|
|
// TODO: check if index.indexPtr will be released by golang, though it occupies little memory
|
|
|
|
// C.free(index.indexPtr)
|
2020-12-18 15:44:27 +08:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewCIndex(typeParams, indexParams map[string]string) (Index, error) {
|
2020-12-21 16:36:07 +08:00
|
|
|
protoTypeParams := &indexcgopb.TypeParams{
|
|
|
|
Params: make([]*commonpb.KeyValuePair, 0),
|
|
|
|
}
|
|
|
|
for key, value := range typeParams {
|
|
|
|
protoTypeParams.Params = append(protoTypeParams.Params, &commonpb.KeyValuePair{Key: key, Value: value})
|
|
|
|
}
|
2020-12-29 14:43:40 +08:00
|
|
|
typeParamsStr := proto.MarshalTextString(protoTypeParams)
|
2020-12-18 15:44:27 +08:00
|
|
|
|
2020-12-21 16:36:07 +08:00
|
|
|
protoIndexParams := &indexcgopb.IndexParams{
|
|
|
|
Params: make([]*commonpb.KeyValuePair, 0),
|
|
|
|
}
|
|
|
|
for key, value := range indexParams {
|
|
|
|
protoIndexParams.Params = append(protoIndexParams.Params, &commonpb.KeyValuePair{Key: key, Value: value})
|
|
|
|
}
|
2020-12-29 14:43:40 +08:00
|
|
|
indexParamsStr := proto.MarshalTextString(protoIndexParams)
|
2020-12-18 15:44:27 +08:00
|
|
|
|
2020-12-29 14:43:40 +08:00
|
|
|
typeParamsPointer := C.CString(typeParamsStr)
|
|
|
|
indexParamsPointer := C.CString(indexParamsStr)
|
2021-02-24 18:25:34 +08:00
|
|
|
defer C.free(unsafe.Pointer(typeParamsPointer))
|
|
|
|
defer C.free(unsafe.Pointer(indexParamsPointer))
|
2020-12-28 16:55:50 +08:00
|
|
|
|
2020-12-21 16:36:07 +08:00
|
|
|
/*
|
2020-12-29 17:53:27 +08:00
|
|
|
CStatus
|
2020-12-28 16:55:50 +08:00
|
|
|
CreateIndex(const char* serialized_type_params,
|
2020-12-29 17:53:27 +08:00
|
|
|
const char* serialized_index_params,
|
|
|
|
CIndex* res_index);
|
2020-12-21 16:36:07 +08:00
|
|
|
*/
|
2020-12-29 17:53:27 +08:00
|
|
|
var indexPtr C.CIndex
|
2021-03-12 19:47:37 +08:00
|
|
|
log.Debug("before create index ...")
|
2020-12-29 17:53:27 +08:00
|
|
|
status := C.CreateIndex(typeParamsPointer, indexParamsPointer, &indexPtr)
|
2021-03-12 19:47:37 +08:00
|
|
|
log.Debug("after create index ...")
|
2020-12-29 17:53:27 +08:00
|
|
|
errorCode := status.error_code
|
|
|
|
if errorCode != 0 {
|
|
|
|
errorMsg := C.GoString(status.error_msg)
|
2021-03-12 19:47:37 +08:00
|
|
|
log.Debug("indexnode", zap.String("create index error msg", errorMsg))
|
2020-12-29 17:53:27 +08:00
|
|
|
defer C.free(unsafe.Pointer(status.error_msg))
|
2021-03-05 10:15:27 +08:00
|
|
|
return nil, fmt.Errorf(" failed, C runtime error detected, error code = %d, err msg = %s", errorCode, errorMsg)
|
2020-12-29 17:53:27 +08:00
|
|
|
}
|
|
|
|
|
2020-12-18 15:44:27 +08:00
|
|
|
return &CIndex{
|
2020-12-29 17:53:27 +08:00
|
|
|
indexPtr: indexPtr,
|
2020-12-18 15:44:27 +08:00
|
|
|
}, nil
|
|
|
|
}
|