From dfc2535aaabf93a302a3c3e288d3ae8a31386353 Mon Sep 17 00:00:00 2001 From: bigsheeper Date: Wed, 15 Sep 2021 10:07:56 +0800 Subject: [PATCH] Use different schema generation functions (#7884) Signed-off-by: bigsheeper --- internal/querynode/impl_test.go | 4 +- internal/querynode/mock_test.go | 51 +++++++++++++---------- internal/querynode/segment_loader_test.go | 4 +- internal/querynode/segment_test.go | 14 +++---- internal/querynode/task_test.go | 6 +-- 5 files changed, 43 insertions(+), 36 deletions(-) diff --git a/internal/querynode/impl_test.go b/internal/querynode/impl_test.go index bd31a1258a..a6ac1f5e56 100644 --- a/internal/querynode/impl_test.go +++ b/internal/querynode/impl_test.go @@ -108,7 +108,7 @@ func TestImpl_WatchDmChannels(t *testing.T) { node, err := genSimpleQueryNode(ctx) assert.NoError(t, err) - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() req := &queryPb.WatchDmChannelsRequest{ Base: &commonpb.MsgBase{ @@ -137,7 +137,7 @@ func TestImpl_LoadSegments(t *testing.T) { node, err := genSimpleQueryNode(ctx) assert.NoError(t, err) - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() req := &queryPb.LoadSegmentsRequest{ Base: &commonpb.MsgBase{ diff --git a/internal/querynode/mock_test.go b/internal/querynode/mock_test.go index 6dbe1ef3d6..efc0e10771 100644 --- a/internal/querynode/mock_test.go +++ b/internal/querynode/mock_test.go @@ -266,13 +266,28 @@ func generateIndex(segmentID UniqueID) ([]string, error) { return indexPaths, nil } -func genSimpleSchema() (*schemapb.CollectionSchema, *schemapb.CollectionSchema) { +func genSimpleSegCoreSchema() *schemapb.CollectionSchema { + fieldVec := genFloatVectorField(simpleVecField) + fieldInt := genConstantField(simpleConstField) + + schema := schemapb.CollectionSchema{ // schema for segCore + Name: defaultCollectionName, + AutoID: true, + Fields: []*schemapb.FieldSchema{ + fieldVec, + fieldInt, + }, + } + return &schema +} + +func genSimpleInsertDataSchema() *schemapb.CollectionSchema { fieldUID := genConstantField(uidField) fieldTimestamp := genConstantField(timestampField) fieldVec := genFloatVectorField(simpleVecField) fieldInt := genConstantField(simpleConstField) - schema1 := schemapb.CollectionSchema{ // schema for insertData + schema := schemapb.CollectionSchema{ // schema for insertData Name: defaultCollectionName, AutoID: true, Fields: []*schemapb.FieldSchema{ @@ -282,15 +297,7 @@ func genSimpleSchema() (*schemapb.CollectionSchema, *schemapb.CollectionSchema) fieldInt, }, } - schema2 := schemapb.CollectionSchema{ // schema for segCore - Name: defaultCollectionName, - AutoID: true, - Fields: []*schemapb.FieldSchema{ - fieldVec, - fieldInt, - }, - } - return &schema1, &schema2 + return &schema } func genCollectionMeta(collectionID UniqueID, schema *schemapb.CollectionSchema) *etcdpb.CollectionMeta { @@ -303,7 +310,7 @@ func genCollectionMeta(collectionID UniqueID, schema *schemapb.CollectionSchema) } func genSimpleCollectionMeta() *etcdpb.CollectionMeta { - simpleSchema, _ := genSimpleSchema() + simpleSchema := genSimpleInsertDataSchema() return genCollectionMeta(defaultCollectionID, simpleSchema) } @@ -389,7 +396,7 @@ func genVectorChunkManager(ctx context.Context) (storage.ChunkManager, error) { } rcm := storage.NewMinioChunkManager(client) - schema, _ := genSimpleSchema() + schema := genSimpleInsertDataSchema() vcm := storage.NewVectorChunkManager(lcm, rcm, &etcdpb.CollectionMeta{ ID: defaultCollectionID, Schema: schema, @@ -501,7 +508,7 @@ func genInsertData(msgLength int, schema *schemapb.CollectionSchema) (*storage.I } func genSimpleInsertData() (*storage.InsertData, error) { - schema, _ := genSimpleSchema() + schema := genSimpleInsertDataSchema() return genInsertData(defaultMsgLength, schema) } @@ -537,7 +544,7 @@ func genStorageBlob(collectionID UniqueID, } func genSimpleStorageBlob() ([]*storage.Blob, error) { - schema, _ := genSimpleSchema() + schema := genSimpleInsertDataSchema() return genStorageBlob(defaultCollectionID, defaultPartitionID, defaultSegmentID, defaultMsgLength, schema) } @@ -599,7 +606,7 @@ func genCommonBlob(msgLength int, schema *schemapb.CollectionSchema) ([]*commonp } func genSimpleCommonBlob() ([]*commonpb.Blob, error) { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() return genCommonBlob(defaultMsgLength, schema) } @@ -648,7 +655,7 @@ func saveBinLog(ctx context.Context, } func saveSimpleBinLog(ctx context.Context) ([]*datapb.FieldBinlog, error) { - schema, _ := genSimpleSchema() + schema := genSimpleInsertDataSchema() return saveBinLog(ctx, defaultCollectionID, defaultPartitionID, defaultSegmentID, defaultMsgLength, schema) } @@ -709,7 +716,7 @@ func genSimpleInsertMsg() (*msgstream.InsertMsg, error) { // ---------- unittest util functions ---------- // functions of replica func genSimpleSealedSegment() (*Segment, error) { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() col := newCollection(defaultCollectionID, schema) seg := newSegment(col, defaultSegmentID, @@ -777,7 +784,7 @@ func genSimpleReplica() (ReplicaInterface, error) { return nil, err } r := newCollectionReplica(kv) - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() err = r.addCollection(defaultCollectionID, schema) if err != nil { return nil, err @@ -887,7 +894,7 @@ func genDSL(schema *schemapb.CollectionSchema, nProb int, topK int64) (string, e } func genSimpleDSL() (string, error) { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() return genDSL(schema, defaultNProb, defaultTopK) } @@ -923,7 +930,7 @@ func genSimplePlaceHolderGroup() ([]byte, error) { } func genSimpleSearchPlanAndRequests() (*SearchPlan, []*searchRequest, error) { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() collection := newCollection(defaultCollectionID, schema) var plan *SearchPlan @@ -1000,7 +1007,7 @@ func genSimpleRetrievePlan() (*RetrievePlan, error) { } timestamp := retrieveMsg.RetrieveRequest.TravelTimestamp - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() collection := newCollection(defaultCollectionID, schema) planExpr, err := genSimpleRetrievePlanExpr() diff --git a/internal/querynode/segment_loader_test.go b/internal/querynode/segment_loader_test.go index f29dd38d2a..6e47a8ebbf 100644 --- a/internal/querynode/segment_loader_test.go +++ b/internal/querynode/segment_loader_test.go @@ -40,7 +40,7 @@ func TestSegmentLoader_loadSegment(t *testing.T) { kv, err := genEtcdKV() assert.NoError(t, err) - schema, _ := genSimpleSchema() + schema := genSimpleInsertDataSchema() fieldBinlog, err := saveSimpleBinLog(ctx) assert.NoError(t, err) @@ -124,7 +124,7 @@ func TestSegmentLoader_notOnService(t *testing.T) { loader := newSegmentLoader(ctx, nil, nil, historical.replica, kv) assert.NotNil(t, loader) - schema, _ := genSimpleSchema() + schema := genSimpleInsertDataSchema() fieldBinlog, err := saveSimpleBinLog(ctx) assert.NoError(t, err) diff --git a/internal/querynode/segment_test.go b/internal/querynode/segment_test.go index c37abac431..a406f13231 100644 --- a/internal/querynode/segment_test.go +++ b/internal/querynode/segment_test.go @@ -819,7 +819,7 @@ func TestSegment_indexInfoTest(t *testing.T) { } func TestSegment_BasicMetrics(t *testing.T) { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() collection := newCollection(defaultCollectionID, schema) segment := newSegment(collection, defaultSegmentID, @@ -878,7 +878,7 @@ func TestSegment_fillVectorFieldsData(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() collection := newCollection(defaultCollectionID, schema) segment := newSegment(collection, defaultSegmentID, @@ -931,7 +931,7 @@ func TestSegment_fillVectorFieldsData(t *testing.T) { } func TestSegment_indexParam(t *testing.T) { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() collection := newCollection(defaultCollectionID, schema) segment := newSegment(collection, defaultSegmentID, @@ -990,7 +990,7 @@ func TestSegment_dropFieldData(t *testing.T) { func TestSegment_updateSegmentIndex(t *testing.T) { t.Run("test updateSegmentIndex invalid", func(t *testing.T) { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() collection := newCollection(defaultCollectionID, schema) segment := newSegment(collection, defaultSegmentID, @@ -1029,7 +1029,7 @@ func TestSegment_updateSegmentIndex(t *testing.T) { func TestSegment_dropSegmentIndex(t *testing.T) { t.Run("test dropSegmentIndex invalid segment type", func(t *testing.T) { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() collection := newCollection(defaultCollectionID, schema) segment := newSegment(collection, defaultSegmentID, @@ -1045,7 +1045,7 @@ func TestSegment_dropSegmentIndex(t *testing.T) { }) t.Run("test dropSegmentIndex nil segment ptr", func(t *testing.T) { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() collection := newCollection(defaultCollectionID, schema) segment := newSegment(collection, defaultSegmentID, @@ -1062,7 +1062,7 @@ func TestSegment_dropSegmentIndex(t *testing.T) { }) t.Run("test dropSegmentIndex nil index", func(t *testing.T) { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() collection := newCollection(defaultCollectionID, schema) segment := newSegment(collection, defaultSegmentID, diff --git a/internal/querynode/task_test.go b/internal/querynode/task_test.go index 67a063752f..0010b46b6a 100644 --- a/internal/querynode/task_test.go +++ b/internal/querynode/task_test.go @@ -27,7 +27,7 @@ func TestTask_watchDmChannelsTask(t *testing.T) { defer cancel() genWatchDMChannelsRequest := func() *querypb.WatchDmChannelsRequest { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() req := &querypb.WatchDmChannelsRequest{ Base: genCommonMsgBase(commonpb.MsgType_WatchDmChannels), CollectionID: defaultCollectionID, @@ -149,7 +149,7 @@ func TestTask_loadSegmentsTask(t *testing.T) { defer cancel() genLoadEmptySegmentsRequest := func() *querypb.LoadSegmentsRequest { - _, schema := genSimpleSchema() + schema := genSimpleSegCoreSchema() req := &querypb.LoadSegmentsRequest{ Base: genCommonMsgBase(commonpb.MsgType_LoadSegments), Schema: schema, @@ -186,7 +186,7 @@ func TestTask_loadSegmentsTask(t *testing.T) { node, err := genSimpleQueryNode(ctx) assert.NoError(t, err) - schema, _ := genSimpleSchema() + schema := genSimpleInsertDataSchema() fieldBinlog, err := saveSimpleBinLog(ctx) assert.NoError(t, err)