milvus/internal/storage/vector_chunk_manager_test.go
XuanYang-cn 7bfde0aa73
[skip e2e]Update license for storage vector chunk m (#14540)
Signed-off-by: yangxuan <xuan.yang@zilliz.com>
2021-12-29 22:42:14 +08:00

279 lines
7.2 KiB
Go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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.
package storage
import (
"context"
"os"
"path"
"strconv"
"testing"
miniokv "github.com/milvus-io/milvus/internal/kv/minio"
"github.com/milvus-io/milvus/internal/proto/etcdpb"
"github.com/milvus-io/milvus/internal/proto/schemapb"
"github.com/milvus-io/milvus/internal/util/paramtable"
"github.com/milvus-io/milvus/internal/util/typeutil"
"github.com/stretchr/testify/assert"
)
func newMinIOKVClient(ctx context.Context, bucketName string) (*miniokv.MinIOKV, error) {
endPoint, _ := Params.Load("_MinioAddress")
accessKeyID, _ := Params.Load("minio.accessKeyID")
secretAccessKey, _ := Params.Load("minio.secretAccessKey")
useSSLStr, _ := Params.Load("minio.useSSL")
useSSL, _ := strconv.ParseBool(useSSLStr)
option := &miniokv.Option{
Address: endPoint,
AccessKeyID: accessKeyID,
SecretAccessKeyID: secretAccessKey,
UseSSL: useSSL,
BucketName: bucketName,
CreateBucket: true,
}
client, err := miniokv.NewMinIOKV(ctx, option)
return client, err
}
func initMeta() *etcdpb.CollectionMeta {
meta := &etcdpb.CollectionMeta{
ID: 1,
CreateTime: 1,
SegmentIDs: []int64{0, 1},
PartitionTags: []string{"partition_0", "partition_1"},
Schema: &schemapb.CollectionSchema{
Name: "schema",
Description: "schema",
AutoID: true,
Fields: []*schemapb.FieldSchema{
{
FieldID: 0,
Name: "row_id",
IsPrimaryKey: false,
Description: "row_id",
DataType: schemapb.DataType_Int64,
},
{
FieldID: 1,
Name: "Ts",
IsPrimaryKey: false,
Description: "Ts",
DataType: schemapb.DataType_Int64,
},
{
FieldID: 101,
Name: "field_int8",
IsPrimaryKey: false,
Description: "description_3",
DataType: schemapb.DataType_Int8,
},
{
FieldID: 108,
Name: "field_binary_vector",
IsPrimaryKey: false,
Description: "description_10",
DataType: schemapb.DataType_BinaryVector,
},
{
FieldID: 109,
Name: "field_float_vector",
IsPrimaryKey: false,
Description: "description_11",
DataType: schemapb.DataType_FloatVector,
},
},
},
}
return meta
}
func initBinlogFile(schema *etcdpb.CollectionMeta) []*Blob {
insertCodec := NewInsertCodec(schema)
insertData := &InsertData{
Data: map[int64]FieldData{
0: &Int64FieldData{
NumRows: []int64{2},
Data: []int64{3, 4},
},
1: &Int64FieldData{
NumRows: []int64{2},
Data: []int64{3, 4},
},
101: &Int8FieldData{
NumRows: []int64{2},
Data: []int8{3, 4},
},
108: &BinaryVectorFieldData{
NumRows: []int64{2},
Data: []byte{0, 255},
Dim: 8,
},
109: &FloatVectorFieldData{
NumRows: []int64{2},
Data: []float32{0, 1, 2, 3, 4, 5, 6, 7, 0, 111, 222, 333, 444, 555, 777, 666},
Dim: 8,
},
},
}
blobs, _, err := insertCodec.Serialize(1, 1, insertData)
if err != nil {
return nil
}
return blobs
}
func buildVectorChunkManager(t *testing.T, localPath string, localCacheEnable bool) (*VectorChunkManager, context.CancelFunc) {
ctx, cancel := context.WithCancel(context.Background())
bucketName := "vector-chunk-manager"
minIOKV, err := newMinIOKVClient(ctx, bucketName)
assert.Nil(t, err)
rcm := NewMinioChunkManager(minIOKV)
lcm := NewLocalChunkManager(localPath)
meta := initMeta()
vcm := NewVectorChunkManager(lcm, rcm, meta, localCacheEnable)
assert.NotNil(t, vcm)
var allCancel context.CancelFunc = func() {
err := minIOKV.RemoveWithPrefix("")
assert.Nil(t, err)
cancel()
}
return vcm, allCancel
}
var Params paramtable.BaseTable
var localPath = "/tmp/milvus/test_data"
func TestMain(m *testing.M) {
Params.Init()
exitCode := m.Run()
err := os.RemoveAll(localPath)
if err != nil {
return
}
os.Exit(exitCode)
}
func TestVectorChunkManager_GetPath(t *testing.T) {
vcm, cancel := buildVectorChunkManager(t, localPath, true)
defer cancel()
assert.NotNil(t, vcm)
key := "1"
err := vcm.Write(key, []byte{1})
assert.Nil(t, err)
pathGet, err := vcm.GetPath(key)
assert.Nil(t, err)
pathJoin := path.Join(localPath, key)
assert.Equal(t, pathGet, pathJoin)
vcm.localCacheEnable = false
err = vcm.remoteChunkManager.Write(key, []byte{1})
assert.Nil(t, err)
pathGet, err = vcm.GetPath(key)
assert.Nil(t, err)
assert.Equal(t, pathGet, key)
}
func TestVectorChunkManager_Write(t *testing.T) {
vcm, cancel := buildVectorChunkManager(t, localPath, false)
defer cancel()
assert.NotNil(t, vcm)
key := "1"
err := vcm.Write(key, []byte{1})
assert.Error(t, err)
vcm.localCacheEnable = true
err = vcm.Write(key, []byte{1})
assert.Nil(t, err)
exist := vcm.Exist(key)
assert.True(t, exist)
}
func TestVectorChunkManager_Read(t *testing.T) {
meta := initMeta()
vcm, cancel := buildVectorChunkManager(t, localPath, false)
defer cancel()
assert.NotNil(t, vcm)
content, err := vcm.Read("9999")
assert.Error(t, err)
assert.Nil(t, content)
vcm.localCacheEnable = true
content, err = vcm.Read("9999")
assert.Error(t, err)
assert.Nil(t, content)
binlogs := initBinlogFile(meta)
assert.NotNil(t, binlogs)
for _, binlog := range binlogs {
err := vcm.remoteChunkManager.Write(binlog.Key, binlog.Value)
assert.Nil(t, err)
}
content, err = vcm.Read("108")
assert.Nil(t, err)
assert.Equal(t, []byte{0, 255}, content)
content, err = vcm.Read("109")
assert.Nil(t, err)
floatResult := make([]float32, 0)
for i := 0; i < len(content)/4; i++ {
singleData := typeutil.BytesToFloat32(content[i*4 : i*4+4])
floatResult = append(floatResult, singleData)
}
assert.Equal(t, []float32{0, 1, 2, 3, 4, 5, 6, 7, 0, 111, 222, 333, 444, 555, 777, 666}, floatResult)
content = make([]byte, 8*4)
byteLen, err := vcm.ReadAt("109", content, 8*4)
assert.Nil(t, err)
assert.Equal(t, 32, byteLen)
floatResult = make([]float32, 0)
for i := 0; i < len(content)/4; i++ {
singleData := typeutil.BytesToFloat32(content[i*4 : i*4+4])
floatResult = append(floatResult, singleData)
}
assert.Equal(t, []float32{0, 111, 222, 333, 444, 555, 777, 666}, floatResult)
byteLen, err = vcm.ReadAt("9999", content, 0)
assert.Error(t, err)
assert.Equal(t, -1, byteLen)
vcm.localCacheEnable = false
byteLen, err = vcm.ReadAt("109", content, 8*4)
assert.Nil(t, err)
assert.Equal(t, 32, byteLen)
byteLen, err = vcm.ReadAt("109", content, 9999)
assert.Error(t, err)
assert.Equal(t, 0, byteLen)
byteLen, err = vcm.ReadAt("9999", content, 0)
assert.Error(t, err)
assert.Equal(t, -1, byteLen)
}