mirror of https://github.com/milvus-io/milvus.git
139 lines
4.0 KiB
Go
139 lines
4.0 KiB
Go
package segcore_test
|
|
|
|
import (
|
|
"context"
|
|
"path/filepath"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
"github.com/milvus-io/milvus/internal/mocks/util/mock_segcore"
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
"github.com/milvus-io/milvus/internal/util/initcore"
|
|
"github.com/milvus-io/milvus/internal/util/segcore"
|
|
"github.com/milvus-io/milvus/pkg/proto/planpb"
|
|
"github.com/milvus-io/milvus/pkg/proto/segcorepb"
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/util/typeutil"
|
|
)
|
|
|
|
func TestGrowingSegment(t *testing.T) {
|
|
paramtable.Init()
|
|
localDataRootPath := filepath.Join(paramtable.Get().LocalStorageCfg.Path.GetValue(), typeutil.QueryNodeRole)
|
|
initcore.InitLocalChunkManager(localDataRootPath)
|
|
err := initcore.InitMmapManager(paramtable.Get())
|
|
assert.NoError(t, err)
|
|
|
|
collectionID := int64(100)
|
|
segmentID := int64(100)
|
|
|
|
schema := mock_segcore.GenTestCollectionSchema("test-reduce", schemapb.DataType_Int64, true)
|
|
collection, err := segcore.CreateCCollection(&segcore.CreateCCollectionRequest{
|
|
CollectionID: collectionID,
|
|
Schema: schema,
|
|
IndexMeta: mock_segcore.GenTestIndexMeta(collectionID, schema),
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, collection)
|
|
defer collection.Release()
|
|
|
|
segment, err := segcore.CreateCSegment(&segcore.CreateCSegmentRequest{
|
|
Collection: collection,
|
|
SegmentID: segmentID,
|
|
SegmentType: segcore.SegmentTypeGrowing,
|
|
IsSorted: false,
|
|
})
|
|
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, segment)
|
|
defer segment.Release()
|
|
|
|
assert.Equal(t, segmentID, segment.ID())
|
|
assert.Equal(t, int64(0), segment.RowNum())
|
|
assert.Zero(t, segment.MemSize())
|
|
assert.True(t, segment.HasRawData(0))
|
|
assertEqualCount(t, collection, segment, 0)
|
|
|
|
insertMsg, err := mock_segcore.GenInsertMsg(collection, 1, segmentID, 100)
|
|
assert.NoError(t, err)
|
|
insertResult, err := segment.Insert(context.Background(), &segcore.InsertRequest{
|
|
RowIDs: insertMsg.RowIDs,
|
|
Timestamps: insertMsg.Timestamps,
|
|
Record: &segcorepb.InsertRecord{
|
|
FieldsData: insertMsg.FieldsData,
|
|
NumRows: int64(len(insertMsg.RowIDs)),
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, insertResult)
|
|
assert.Equal(t, int64(100), insertResult.InsertedRows)
|
|
|
|
assert.Equal(t, int64(100), segment.RowNum())
|
|
assertEqualCount(t, collection, segment, 100)
|
|
|
|
pk := storage.NewInt64PrimaryKeys(1)
|
|
pk.Append(storage.NewInt64PrimaryKey(10))
|
|
deleteResult, err := segment.Delete(context.Background(), &segcore.DeleteRequest{
|
|
PrimaryKeys: pk,
|
|
Timestamps: []typeutil.Timestamp{
|
|
1000,
|
|
},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, deleteResult)
|
|
|
|
assert.Equal(t, int64(99), segment.RowNum())
|
|
}
|
|
|
|
func assertEqualCount(
|
|
t *testing.T,
|
|
collection *segcore.CCollection,
|
|
segment segcore.CSegment,
|
|
count int64,
|
|
) {
|
|
plan := planpb.PlanNode{
|
|
Node: &planpb.PlanNode_Query{
|
|
Query: &planpb.QueryPlanNode{
|
|
IsCount: true,
|
|
},
|
|
},
|
|
}
|
|
expr, err := proto.Marshal(&plan)
|
|
assert.NoError(t, err)
|
|
retrievePlan, err := segcore.NewRetrievePlan(
|
|
collection,
|
|
expr,
|
|
typeutil.MaxTimestamp,
|
|
100,
|
|
0)
|
|
defer retrievePlan.Delete()
|
|
|
|
assert.True(t, retrievePlan.ShouldIgnoreNonPk())
|
|
assert.False(t, retrievePlan.IsIgnoreNonPk())
|
|
retrievePlan.SetIgnoreNonPk(true)
|
|
assert.True(t, retrievePlan.IsIgnoreNonPk())
|
|
assert.NotZero(t, retrievePlan.MsgID())
|
|
|
|
assert.NotNil(t, retrievePlan)
|
|
assert.NoError(t, err)
|
|
|
|
retrieveResult, err := segment.Retrieve(context.Background(), retrievePlan)
|
|
assert.NotNil(t, retrieveResult)
|
|
assert.NoError(t, err)
|
|
result, err := retrieveResult.GetResult()
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
assert.Equal(t, count, result.AllRetrieveCount)
|
|
retrieveResult.Release()
|
|
|
|
retrieveResult2, err := segment.RetrieveByOffsets(context.Background(), &segcore.RetrievePlanWithOffsets{
|
|
RetrievePlan: retrievePlan,
|
|
Offsets: []int64{0, 1, 2, 3, 4},
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, retrieveResult2)
|
|
retrieveResult2.Release()
|
|
}
|