wehub-resource-sync
498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
703 lines
20 KiB
Go
703 lines
20 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 segments
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/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/internal/util/streamrpc"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
|
)
|
|
|
|
type RetrieveSuite struct {
|
|
suite.Suite
|
|
|
|
// schema
|
|
ctx context.Context
|
|
schema *schemapb.CollectionSchema
|
|
|
|
// Dependencies
|
|
rootPath string
|
|
chunkManager storage.ChunkManager
|
|
|
|
// Data
|
|
manager *Manager
|
|
collectionID int64
|
|
partitionID int64
|
|
segmentID int64
|
|
collection *Collection
|
|
sealed Segment
|
|
growing Segment
|
|
}
|
|
|
|
func (suite *RetrieveSuite) SetupSuite() {
|
|
paramtable.Init()
|
|
}
|
|
|
|
func (suite *RetrieveSuite) SetupTest() {
|
|
var err error
|
|
suite.ctx = context.Background()
|
|
msgLength := 100
|
|
|
|
suite.rootPath = suite.T().Name()
|
|
chunkManagerFactory := storage.NewTestChunkManagerFactory(paramtable.Get(), suite.rootPath)
|
|
suite.chunkManager, _ = chunkManagerFactory.NewPersistentStorageChunkManager(suite.ctx)
|
|
initcore.InitRemoteChunkManager(paramtable.Get())
|
|
initcore.InitLocalChunkManager(suite.rootPath)
|
|
initcore.InitMmapManager(paramtable.Get(), 1)
|
|
initcore.InitTieredStorage(paramtable.Get())
|
|
|
|
suite.collectionID = 100
|
|
suite.partitionID = 10
|
|
suite.segmentID = 100
|
|
|
|
suite.manager = NewManager()
|
|
suite.schema = mock_segcore.GenTestCollectionSchema("test-reduce", schemapb.DataType_Int64, true)
|
|
indexMeta := mock_segcore.GenTestIndexMeta(suite.collectionID, suite.schema)
|
|
suite.manager.Collection.PutOrRef(suite.collectionID,
|
|
suite.schema,
|
|
indexMeta,
|
|
&querypb.LoadMetaInfo{
|
|
LoadType: querypb.LoadType_LoadCollection,
|
|
CollectionID: suite.collectionID,
|
|
PartitionIDs: []int64{suite.partitionID},
|
|
},
|
|
)
|
|
suite.collection = suite.manager.Collection.Get(suite.collectionID)
|
|
loader := NewLoader(suite.ctx, suite.manager, suite.chunkManager)
|
|
|
|
binlogs, statslogs, err := mock_segcore.SaveBinLog(suite.ctx,
|
|
suite.collectionID,
|
|
suite.partitionID,
|
|
suite.segmentID,
|
|
msgLength,
|
|
suite.schema,
|
|
suite.chunkManager,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
sealLoadInfo := querypb.SegmentLoadInfo{
|
|
SegmentID: suite.segmentID,
|
|
CollectionID: suite.collectionID,
|
|
PartitionID: suite.partitionID,
|
|
NumOfRows: int64(msgLength),
|
|
BinlogPaths: binlogs,
|
|
Statslogs: statslogs,
|
|
InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", suite.collectionID),
|
|
Level: datapb.SegmentLevel_Legacy,
|
|
}
|
|
|
|
suite.sealed, err = NewSegment(suite.ctx,
|
|
suite.collection,
|
|
suite.manager.Segment,
|
|
SegmentTypeSealed,
|
|
0,
|
|
&sealLoadInfo,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
bfs, err := loader.loadSingleBloomFilterSet(suite.ctx, suite.collectionID, &sealLoadInfo, SegmentTypeSealed)
|
|
suite.Require().NoError(err)
|
|
suite.sealed.SetPKCandidate(bfs)
|
|
|
|
for _, binlog := range binlogs {
|
|
err = suite.sealed.(*LocalSegment).LoadFieldData(suite.ctx, binlog.FieldID, int64(msgLength), binlog)
|
|
suite.Require().NoError(err)
|
|
}
|
|
|
|
binlogs, statlogs, err := mock_segcore.SaveBinLog(suite.ctx,
|
|
suite.collectionID,
|
|
suite.partitionID,
|
|
suite.segmentID+1,
|
|
msgLength,
|
|
suite.schema,
|
|
suite.chunkManager,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
growingLoadInfo := querypb.SegmentLoadInfo{
|
|
SegmentID: suite.segmentID + 1,
|
|
CollectionID: suite.collectionID,
|
|
PartitionID: suite.partitionID,
|
|
BinlogPaths: binlogs,
|
|
Statslogs: statlogs,
|
|
InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", suite.collectionID),
|
|
Level: datapb.SegmentLevel_Legacy,
|
|
}
|
|
|
|
// allow growing segment use the bloom filter
|
|
paramtable.Get().QueryNodeCfg.SkipGrowingSegmentBF.SwapTempValue("false")
|
|
|
|
suite.growing, err = NewSegment(suite.ctx,
|
|
suite.collection,
|
|
suite.manager.Segment,
|
|
SegmentTypeGrowing,
|
|
0,
|
|
&growingLoadInfo,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
bfs, err = loader.loadSingleBloomFilterSet(suite.ctx, suite.collectionID, &growingLoadInfo, SegmentTypeGrowing)
|
|
suite.Require().NoError(err)
|
|
suite.growing.SetPKCandidate(bfs)
|
|
|
|
insertMsg, err := mock_segcore.GenInsertMsg(suite.collection.GetCCollection(), suite.partitionID, suite.growing.ID(), msgLength)
|
|
suite.Require().NoError(err)
|
|
insertRecord, _, err := storage.TransferInsertMsgToInsertRecord(suite.collection.Schema(), insertMsg)
|
|
suite.Require().NoError(err)
|
|
err = suite.growing.Insert(suite.ctx, insertMsg.RowIDs, insertMsg.Timestamps, insertRecord)
|
|
suite.Require().NoError(err)
|
|
|
|
suite.manager.Segment.Put(context.Background(), SegmentTypeSealed, suite.sealed)
|
|
suite.manager.Segment.Put(context.Background(), SegmentTypeGrowing, suite.growing)
|
|
}
|
|
|
|
func (suite *RetrieveSuite) TearDownTest() {
|
|
suite.sealed.Release(context.Background())
|
|
suite.growing.Release(context.Background())
|
|
DeleteCollection(suite.collection)
|
|
ctx := context.Background()
|
|
suite.chunkManager.RemoveWithPrefix(ctx, suite.rootPath)
|
|
}
|
|
|
|
func (suite *RetrieveSuite) TestRetrieveSealed() {
|
|
plan, err := mock_segcore.GenSimpleRetrievePlan(suite.collection.GetCCollection())
|
|
suite.NoError(err)
|
|
|
|
req := &querypb.QueryRequest{
|
|
Req: &internalpb.RetrieveRequest{
|
|
PkFilter: 0,
|
|
CollectionID: suite.collectionID,
|
|
PartitionIDs: []int64{suite.partitionID},
|
|
},
|
|
SegmentIDs: []int64{suite.sealed.ID()},
|
|
Scope: querypb.DataScope_Historical,
|
|
}
|
|
|
|
res, segments, err := Retrieve(context.TODO(), suite.manager, plan, req)
|
|
suite.NoError(err)
|
|
suite.Len(res[0].Result.Offset, 3)
|
|
suite.manager.Segment.Unpin(segments)
|
|
|
|
resultByOffsets, err := suite.sealed.RetrieveByOffsets(context.Background(), &segcore.RetrievePlanWithOffsets{
|
|
RetrievePlan: plan,
|
|
Offsets: []int64{0, 1},
|
|
})
|
|
suite.NoError(err)
|
|
suite.Len(resultByOffsets.Offset, 0)
|
|
}
|
|
|
|
func (suite *RetrieveSuite) TestRetrieveWithFilter() {
|
|
plan, err := mock_segcore.GenSimpleRetrievePlan(suite.collection.GetCCollection())
|
|
suite.NoError(err)
|
|
|
|
suite.Run("SegmentFilterRules", func() {
|
|
// create more 10 seal segments to test BF
|
|
// The pk in seg range is {segN [0...N-1]}
|
|
// ex.
|
|
// seg1 [0]
|
|
// seg2 [0, 1]
|
|
// ...
|
|
// seg10 [0, 1, 2, ..., 9]
|
|
loader := NewLoader(suite.ctx, suite.manager, suite.chunkManager)
|
|
for i := range 10 {
|
|
segid := int64(i + 1)
|
|
msgLen := i + 1
|
|
bl, sl, err := mock_segcore.SaveBinLog(suite.ctx,
|
|
suite.collectionID,
|
|
suite.partitionID,
|
|
segid,
|
|
msgLen,
|
|
suite.schema,
|
|
suite.chunkManager)
|
|
|
|
suite.Require().NoError(err)
|
|
|
|
sealLoadInfo := querypb.SegmentLoadInfo{
|
|
SegmentID: segid,
|
|
CollectionID: suite.collectionID,
|
|
PartitionID: suite.partitionID,
|
|
NumOfRows: int64(msgLen),
|
|
BinlogPaths: bl,
|
|
Statslogs: sl,
|
|
InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", suite.collectionID),
|
|
Level: datapb.SegmentLevel_Legacy,
|
|
}
|
|
|
|
sealseg, err := NewSegment(suite.ctx,
|
|
suite.collection,
|
|
suite.manager.Segment,
|
|
SegmentTypeSealed,
|
|
0,
|
|
&sealLoadInfo,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
bfs, err := loader.loadSingleBloomFilterSet(suite.ctx, suite.collectionID,
|
|
&sealLoadInfo, SegmentTypeSealed)
|
|
suite.Require().NoError(err)
|
|
sealseg.SetPKCandidate(bfs)
|
|
|
|
suite.manager.Segment.Put(suite.ctx, SegmentTypeSealed, sealseg)
|
|
}
|
|
|
|
req := &querypb.QueryRequest{
|
|
Req: &internalpb.RetrieveRequest{
|
|
PkFilter: 0,
|
|
CollectionID: suite.collectionID,
|
|
PartitionIDs: []int64{suite.partitionID},
|
|
},
|
|
SegmentIDs: []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
|
|
Scope: querypb.DataScope_Historical,
|
|
}
|
|
|
|
// PK filtering now happens at the delegator layer, not at the worker layer.
|
|
// All requested segments are returned regardless of the expression.
|
|
res, segments, err := Retrieve(context.TODO(), suite.manager, plan, req)
|
|
suite.NoError(err)
|
|
suite.Len(res, 10)
|
|
suite.manager.Segment.Unpin(segments)
|
|
|
|
// remove the segs
|
|
for i := range 10 {
|
|
suite.manager.Segment.Remove(suite.ctx, int64(i+1) /*segmentID*/, querypb.DataScope_Historical)
|
|
}
|
|
})
|
|
}
|
|
|
|
func (suite *RetrieveSuite) TestRetrieveGrowing() {
|
|
plan, err := mock_segcore.GenSimpleRetrievePlan(suite.collection.GetCCollection())
|
|
suite.NoError(err)
|
|
|
|
req := &querypb.QueryRequest{
|
|
Req: &internalpb.RetrieveRequest{
|
|
PkFilter: 0,
|
|
CollectionID: suite.collectionID,
|
|
PartitionIDs: []int64{suite.partitionID},
|
|
},
|
|
SegmentIDs: []int64{suite.growing.ID()},
|
|
Scope: querypb.DataScope_Streaming,
|
|
}
|
|
|
|
res, segments, err := Retrieve(context.TODO(), suite.manager, plan, req)
|
|
suite.NoError(err)
|
|
suite.Len(res[0].Result.Offset, 3)
|
|
suite.manager.Segment.Unpin(segments)
|
|
|
|
resultByOffsets, err := suite.growing.RetrieveByOffsets(context.Background(), &segcore.RetrievePlanWithOffsets{
|
|
RetrievePlan: plan,
|
|
Offsets: []int64{0, 1},
|
|
})
|
|
suite.NoError(err)
|
|
suite.Len(resultByOffsets.Offset, 0)
|
|
}
|
|
|
|
func (suite *RetrieveSuite) TestRetrieveStreamWithFilterDoesNotPruneGrowing() {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
loader := NewLoader(ctx, suite.manager, suite.chunkManager)
|
|
growingSegIDs := make([]int64, 0, 10)
|
|
for i := range 10 {
|
|
segID := int64(i + 3000)
|
|
msgLen := i + 1
|
|
binlogs, statslogs, err := mock_segcore.SaveBinLog(ctx,
|
|
suite.collectionID,
|
|
suite.partitionID,
|
|
segID,
|
|
msgLen,
|
|
suite.schema,
|
|
suite.chunkManager,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
loadInfo := &querypb.SegmentLoadInfo{
|
|
SegmentID: segID,
|
|
CollectionID: suite.collectionID,
|
|
PartitionID: suite.partitionID,
|
|
BinlogPaths: binlogs,
|
|
Statslogs: statslogs,
|
|
InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", suite.collectionID),
|
|
Level: datapb.SegmentLevel_Legacy,
|
|
}
|
|
|
|
seg, err := NewSegment(ctx,
|
|
suite.collection,
|
|
suite.manager.Segment,
|
|
SegmentTypeGrowing,
|
|
0,
|
|
loadInfo,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
bfs, err := loader.loadSingleBloomFilterSet(ctx, suite.collectionID, loadInfo, SegmentTypeGrowing)
|
|
suite.Require().NoError(err)
|
|
seg.SetPKCandidate(bfs)
|
|
|
|
insertMsg, err := mock_segcore.GenInsertMsg(suite.collection.GetCCollection(), suite.partitionID, segID, msgLen)
|
|
suite.Require().NoError(err)
|
|
insertRecord, _, err := storage.TransferInsertMsgToInsertRecord(suite.collection.Schema(), insertMsg)
|
|
suite.Require().NoError(err)
|
|
err = seg.Insert(ctx, insertMsg.RowIDs, insertMsg.Timestamps, insertRecord)
|
|
suite.Require().NoError(err)
|
|
|
|
suite.manager.Segment.Put(ctx, SegmentTypeGrowing, seg)
|
|
growingSegIDs = append(growingSegIDs, segID)
|
|
}
|
|
|
|
plan, err := mock_segcore.GenSimpleRetrievePlan(suite.collection.GetCCollection())
|
|
suite.NoError(err)
|
|
|
|
req := &querypb.QueryRequest{
|
|
Req: &internalpb.RetrieveRequest{
|
|
CollectionID: suite.collectionID,
|
|
PartitionIDs: []int64{suite.partitionID},
|
|
},
|
|
SegmentIDs: growingSegIDs,
|
|
Scope: querypb.DataScope_Streaming,
|
|
}
|
|
|
|
client := streamrpc.NewLocalQueryClient(ctx)
|
|
server := client.CreateServer()
|
|
|
|
var retrievedSegments []Segment
|
|
go func() {
|
|
retrievedSegments, err = RetrieveStream(ctx, suite.manager, plan, req, server)
|
|
suite.NoError(err)
|
|
server.FinishSend(err)
|
|
}()
|
|
|
|
for {
|
|
_, err := client.Recv()
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
}
|
|
|
|
suite.Len(retrievedSegments, 10)
|
|
suite.manager.Segment.Unpin(retrievedSegments)
|
|
|
|
for _, segID := range growingSegIDs {
|
|
suite.manager.Segment.Remove(ctx, segID, querypb.DataScope_Streaming)
|
|
}
|
|
}
|
|
|
|
func (suite *RetrieveSuite) TestRetrieveWithFilterDoesNotPruneGrowing() {
|
|
ctx := context.Background()
|
|
loader := NewLoader(ctx, suite.manager, suite.chunkManager)
|
|
|
|
growingSegIDs := make([]int64, 0, 10)
|
|
for i := range 10 {
|
|
segID := int64(i + 3000)
|
|
msgLen := i + 1
|
|
binlogs, statslogs, err := mock_segcore.SaveBinLog(ctx,
|
|
suite.collectionID,
|
|
suite.partitionID,
|
|
segID,
|
|
msgLen,
|
|
suite.schema,
|
|
suite.chunkManager,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
loadInfo := &querypb.SegmentLoadInfo{
|
|
SegmentID: segID,
|
|
CollectionID: suite.collectionID,
|
|
PartitionID: suite.partitionID,
|
|
BinlogPaths: binlogs,
|
|
Statslogs: statslogs,
|
|
InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", suite.collectionID),
|
|
Level: datapb.SegmentLevel_Legacy,
|
|
}
|
|
|
|
seg, err := NewSegment(ctx,
|
|
suite.collection,
|
|
suite.manager.Segment,
|
|
SegmentTypeGrowing,
|
|
0,
|
|
loadInfo,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
bfs, err := loader.loadSingleBloomFilterSet(ctx, suite.collectionID, loadInfo, SegmentTypeGrowing)
|
|
suite.Require().NoError(err)
|
|
seg.SetPKCandidate(bfs)
|
|
|
|
insertMsg, err := mock_segcore.GenInsertMsg(suite.collection.GetCCollection(), suite.partitionID, segID, msgLen)
|
|
suite.Require().NoError(err)
|
|
insertRecord, _, err := storage.TransferInsertMsgToInsertRecord(suite.collection.Schema(), insertMsg)
|
|
suite.Require().NoError(err)
|
|
err = seg.Insert(ctx, insertMsg.RowIDs, insertMsg.Timestamps, insertRecord)
|
|
suite.Require().NoError(err)
|
|
|
|
suite.manager.Segment.Put(ctx, SegmentTypeGrowing, seg)
|
|
growingSegIDs = append(growingSegIDs, segID)
|
|
}
|
|
|
|
plan, err := mock_segcore.GenSimpleRetrievePlan(suite.collection.GetCCollection())
|
|
suite.NoError(err)
|
|
|
|
req := &querypb.QueryRequest{
|
|
Req: &internalpb.RetrieveRequest{
|
|
CollectionID: suite.collectionID,
|
|
PartitionIDs: []int64{suite.partitionID},
|
|
},
|
|
SegmentIDs: growingSegIDs,
|
|
Scope: querypb.DataScope_Streaming,
|
|
}
|
|
|
|
res, segments, err := Retrieve(ctx, suite.manager, plan, req)
|
|
suite.NoError(err)
|
|
suite.Len(res, 10)
|
|
suite.Len(segments, 10)
|
|
suite.manager.Segment.Unpin(segments)
|
|
|
|
for _, segID := range growingSegIDs {
|
|
suite.manager.Segment.Remove(ctx, segID, querypb.DataScope_Streaming)
|
|
}
|
|
}
|
|
|
|
func (suite *RetrieveSuite) TestRetrieveStreamSealed() {
|
|
plan, err := mock_segcore.GenSimpleRetrievePlan(suite.collection.GetCCollection())
|
|
suite.NoError(err)
|
|
|
|
req := &querypb.QueryRequest{
|
|
Req: &internalpb.RetrieveRequest{
|
|
PkFilter: 0,
|
|
CollectionID: suite.collectionID,
|
|
PartitionIDs: []int64{suite.partitionID},
|
|
},
|
|
SegmentIDs: []int64{suite.sealed.ID()},
|
|
Scope: querypb.DataScope_Historical,
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
client := streamrpc.NewLocalQueryClient(ctx)
|
|
server := client.CreateServer()
|
|
|
|
go func() {
|
|
segments, err := RetrieveStream(ctx, suite.manager, plan, req, server)
|
|
suite.NoError(err)
|
|
suite.manager.Segment.Unpin(segments)
|
|
server.FinishSend(err)
|
|
}()
|
|
|
|
sum := 0
|
|
for {
|
|
result, err := client.Recv()
|
|
if err != nil {
|
|
if err == io.EOF {
|
|
suite.Equal(3, sum)
|
|
break
|
|
}
|
|
suite.Fail("Retrieve stream fetch error")
|
|
}
|
|
|
|
err = merr.Error(result.GetStatus())
|
|
suite.NoError(err)
|
|
|
|
sum += len(result.Ids.GetIntId().GetData())
|
|
}
|
|
}
|
|
|
|
func (suite *RetrieveSuite) TestRetrieveStreamWithFilter() {
|
|
plan, err := mock_segcore.GenSimpleRetrievePlan(suite.collection.GetCCollection())
|
|
suite.NoError(err)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
// create more sealed segments for testing
|
|
loader := NewLoader(ctx, suite.manager, suite.chunkManager)
|
|
for i := range 10 {
|
|
segID := int64(i + 2000)
|
|
msgLen := i + 1
|
|
binlogs, statslogs, err := mock_segcore.SaveBinLog(ctx,
|
|
suite.collectionID,
|
|
suite.partitionID,
|
|
segID,
|
|
msgLen,
|
|
suite.schema,
|
|
suite.chunkManager,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
loadInfo := &querypb.SegmentLoadInfo{
|
|
SegmentID: segID,
|
|
CollectionID: suite.collectionID,
|
|
PartitionID: suite.partitionID,
|
|
NumOfRows: int64(msgLen),
|
|
BinlogPaths: binlogs,
|
|
Statslogs: statslogs,
|
|
InsertChannel: fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", suite.collectionID),
|
|
Level: datapb.SegmentLevel_Legacy,
|
|
}
|
|
|
|
seg, err := NewSegment(ctx,
|
|
suite.collection,
|
|
suite.manager.Segment,
|
|
SegmentTypeSealed,
|
|
0,
|
|
loadInfo,
|
|
)
|
|
suite.Require().NoError(err)
|
|
|
|
bfs, err := loader.loadSingleBloomFilterSet(ctx, suite.collectionID, loadInfo, SegmentTypeSealed)
|
|
suite.Require().NoError(err)
|
|
seg.SetPKCandidate(bfs)
|
|
|
|
for _, binlog := range binlogs {
|
|
err = seg.(*LocalSegment).LoadFieldData(ctx, binlog.FieldID, int64(msgLen), binlog)
|
|
suite.Require().NoError(err)
|
|
}
|
|
|
|
suite.manager.Segment.Put(ctx, SegmentTypeSealed, seg)
|
|
}
|
|
|
|
segIDs := []int64{2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009}
|
|
|
|
suite.Run("WithSegmentFilter", func() {
|
|
req := &querypb.QueryRequest{
|
|
Req: &internalpb.RetrieveRequest{
|
|
PkFilter: 0,
|
|
CollectionID: suite.collectionID,
|
|
PartitionIDs: []int64{suite.partitionID},
|
|
},
|
|
SegmentIDs: segIDs,
|
|
Scope: querypb.DataScope_Historical,
|
|
}
|
|
|
|
client := streamrpc.NewLocalQueryClient(ctx)
|
|
server := client.CreateServer()
|
|
|
|
var retrievedSegments []Segment
|
|
go func() {
|
|
retrievedSegments, err = RetrieveStream(ctx, suite.manager, plan, req, server)
|
|
suite.NoError(err)
|
|
server.FinishSend(err)
|
|
}()
|
|
|
|
// consume all results
|
|
for {
|
|
_, err := client.Recv()
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
}
|
|
|
|
// PK filtering now happens at the delegator layer, not the worker layer.
|
|
// All 10 requested segments should be retrieved.
|
|
suite.Len(retrievedSegments, 10)
|
|
suite.manager.Segment.Unpin(retrievedSegments)
|
|
})
|
|
|
|
// cleanup
|
|
for _, segID := range segIDs {
|
|
suite.manager.Segment.Remove(ctx, segID, querypb.DataScope_Historical)
|
|
}
|
|
}
|
|
|
|
func (suite *RetrieveSuite) TestRetrieveNonExistSegment() {
|
|
plan, err := mock_segcore.GenSimpleRetrievePlan(suite.collection.GetCCollection())
|
|
suite.NoError(err)
|
|
|
|
req := &querypb.QueryRequest{
|
|
Req: &internalpb.RetrieveRequest{
|
|
PkFilter: 0,
|
|
CollectionID: suite.collectionID,
|
|
PartitionIDs: []int64{suite.partitionID},
|
|
},
|
|
SegmentIDs: []int64{999},
|
|
Scope: querypb.DataScope_Streaming,
|
|
}
|
|
|
|
res, segments, err := Retrieve(context.TODO(), suite.manager, plan, req)
|
|
suite.Error(err)
|
|
suite.Len(res, 0)
|
|
suite.manager.Segment.Unpin(segments)
|
|
}
|
|
|
|
func (suite *RetrieveSuite) TestRetrieveNilSegment() {
|
|
plan, err := mock_segcore.GenSimpleRetrievePlan(suite.collection.GetCCollection())
|
|
suite.NoError(err)
|
|
|
|
suite.sealed.Release(context.Background())
|
|
req := &querypb.QueryRequest{
|
|
Req: &internalpb.RetrieveRequest{
|
|
PkFilter: 0,
|
|
CollectionID: suite.collectionID,
|
|
PartitionIDs: []int64{suite.partitionID},
|
|
},
|
|
SegmentIDs: []int64{suite.sealed.ID()},
|
|
Scope: querypb.DataScope_Historical,
|
|
}
|
|
|
|
res, segments, err := Retrieve(context.TODO(), suite.manager, plan, req)
|
|
suite.ErrorIs(err, merr.ErrSegmentNotLoaded)
|
|
suite.Len(res, 0)
|
|
suite.manager.Segment.Unpin(segments)
|
|
}
|
|
|
|
func TestRetrieve(t *testing.T) {
|
|
suite.Run(t, new(RetrieveSuite))
|
|
}
|
|
|
|
func TestShouldEnableIgnoreNonPkWithGroupBy(t *testing.T) {
|
|
t.Run("plain multi segment", func(t *testing.T) {
|
|
req := &querypb.QueryRequest{Req: &internalpb.RetrieveRequest{
|
|
Limit: 10,
|
|
}}
|
|
assert.True(t, shouldEnableIgnoreNonPk(req, 3, true))
|
|
})
|
|
|
|
t.Run("group by multi segment", func(t *testing.T) {
|
|
req := &querypb.QueryRequest{Req: &internalpb.RetrieveRequest{
|
|
Limit: 10,
|
|
GroupByFieldIds: []int64{100},
|
|
Aggregates: []*planpb.Aggregate{
|
|
{Op: planpb.AggregateOp_count, FieldId: 500},
|
|
},
|
|
}}
|
|
assert.False(t, shouldEnableIgnoreNonPk(req, 3, true))
|
|
})
|
|
|
|
t.Run("single segment", func(t *testing.T) {
|
|
req := &querypb.QueryRequest{Req: &internalpb.RetrieveRequest{
|
|
Limit: 10,
|
|
}}
|
|
assert.False(t, shouldEnableIgnoreNonPk(req, 1, true))
|
|
})
|
|
}
|