milvus/internal/querynode/search_service_test.go

package querynode

import (
	"context"
	"encoding/binary"
	"math"
	"math/rand"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/stretchr/testify/assert"

	"github.com/zilliztech/milvus-distributed/internal/msgstream"
	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb"
	"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
)

func loadFields(segment *Segment, DIM int, N int) error {
	// generate vector field
	vectorFieldID := int64(100)
	vectors := make([]float32, N*DIM)
	for i := 0; i < N*DIM; i++ {
		vectors[i] = rand.Float32()
	}

	// generate int field
	agesFieldID := int64(101)
	ages := make([]int32, N)
	for i := 0; i < N; i++ {
		ages[i] = int32(N)
	}

	err := segment.segmentLoadFieldData(vectorFieldID, N, vectors)
	if err != nil {
		return err
	}
	err = segment.segmentLoadFieldData(agesFieldID, N, ages)
	if err != nil {
		return err
	}
	rowIDs := ages
	err = segment.segmentLoadFieldData(rowIDFieldID, N, rowIDs)
	return err
}

func sendSearchRequest(ctx context.Context, DIM int) error {
	// init message stream
	msFactory, err := newMessageStreamFactory()
	if err != nil {
		return err
	}
	searchProducerChannels := Params.SearchChannelNames

	searchStream, _ := msFactory.NewMsgStream(ctx)
	searchStream.AsProducer(searchProducerChannels)
	searchStream.Start()

	// generate search rawData
	var vec = make([]float32, DIM)
	for i := 0; i < DIM; i++ {
		vec[i] = rand.Float32()
	}
	dslString := "{\"bool\": { \n\"vector\": {\n \"vec\": {\n \"metric_type\": \"L2\", \n \"params\": {\n \"nprobe\": 10 \n},\n \"query\": \"$0\",\"topk\": 10 \n } \n } \n } \n }"
	var searchRawData1 []byte
	var searchRawData2 []byte
	for i, ele := range vec {
		buf := make([]byte, 4)
		binary.LittleEndian.PutUint32(buf, math.Float32bits(ele+float32(i*2)))
		searchRawData1 = append(searchRawData1, buf...)
	}
	for i, ele := range vec {
		buf := make([]byte, 4)
		binary.LittleEndian.PutUint32(buf, math.Float32bits(ele+float32(i*4)))
		searchRawData2 = append(searchRawData2, buf...)
	}

	// generate placeholder
	placeholderValue := milvuspb.PlaceholderValue{
		Tag:    "$0",
		Type:   milvuspb.PlaceholderType_FloatVector,
		Values: [][]byte{searchRawData1, searchRawData2},
	}
	placeholderGroup := milvuspb.PlaceholderGroup{
		Placeholders: []*milvuspb.PlaceholderValue{&placeholderValue},
	}
	placeGroupByte, err := proto.Marshal(&placeholderGroup)
	if err != nil {
		return err
	}

	// generate searchRequest
	searchReq := milvuspb.SearchRequest{
		Dsl:              dslString,
		PlaceholderGroup: placeGroupByte,
	}
	searchReqBytes, err := proto.Marshal(&searchReq)
	if err != nil {
		return err
	}
	blob := commonpb.Blob{
		Value: searchReqBytes,
	}

	// generate searchMsg
	searchMsg := &msgstream.SearchMsg{
		BaseMsg: msgstream.BaseMsg{
			HashValues: []uint32{0},
		},
		SearchRequest: internalpb.SearchRequest{
			Base: &commonpb.MsgBase{
				MsgType:   commonpb.MsgType_Search,
				MsgID:     1,
				Timestamp: Timestamp(10),
				SourceID:  1,
			},
			ResultChannelID: "0",
			Query:           &blob,
		},
	}
	msgPackSearch := msgstream.MsgPack{}
	msgPackSearch.Msgs = append(msgPackSearch.Msgs, searchMsg)

	// produce search message
	err = searchStream.Produce(&msgPackSearch)
	return err
}

func sendTimeTick(ctx context.Context) error {
	// init message stream
	msFactory, err := newMessageStreamFactory()
	if err != nil {
		return err
	}

	// generate timeTick
	timeTickMsgPack := msgstream.MsgPack{}
	baseMsg := msgstream.BaseMsg{
		BeginTimestamp: Timestamp(20),
		EndTimestamp:   Timestamp(20),
		HashValues:     []uint32{0},
	}
	timeTickResult := internalpb.TimeTickMsg{
		Base: &commonpb.MsgBase{
			MsgType:   commonpb.MsgType_TimeTick,
			MsgID:     0,
			Timestamp: Timestamp(20),
			SourceID:  0,
		},
	}
	timeTickMsg := &msgstream.TimeTickMsg{
		BaseMsg:     baseMsg,
		TimeTickMsg: timeTickResult,
	}
	timeTickMsgPack.Msgs = append(timeTickMsgPack.Msgs, timeTickMsg)

	// produce timeTick message
	insertChannels := Params.InsertChannelNames
	insertStream, _ := msFactory.NewMsgStream(ctx)
	insertStream.AsProducer(insertChannels)
	insertStream.Start()

	err = insertStream.Broadcast(&timeTickMsgPack)
	return err
}

func TestSearch_Search(t *testing.T) {
	const N = 10000
	const DIM = 16

	// init queryNode
	collectionID := UniqueID(0)
	segmentID := UniqueID(1)
	node := newQueryNodeMock()
	initTestMeta(t, node, collectionID, UniqueID(0))

	msFactory, err := newMessageStreamFactory()
	assert.NoError(t, err)

	// start dataSync
	newDS := newDataSyncService(node.queryNodeLoopCtx, node.replica, msFactory, collectionID)
	err = node.addDataSyncService(collectionID, newDS)
	assert.NoError(t, err)
	ds, err := node.getDataSyncService(collectionID)
	assert.NoError(t, err)
	go ds.start()

	// start search service
	node.searchService = newSearchService(node.queryNodeLoopCtx, node.replica, msFactory)
	go node.searchService.start()
	node.searchService.startSearchCollection(collectionID)

	tSafe := node.replica.getTSafe(collectionID)
	assert.NotNil(t, tSafe)
	tSafe.set(1000)

	// load segment
	err = node.replica.addSegment(segmentID, defaultPartitionID, collectionID, segmentTypeSealed)
	assert.NoError(t, err)
	segment, err := node.replica.getSegmentByID(segmentID)
	assert.NoError(t, err)
	err = loadFields(segment, DIM, N)
	assert.NoError(t, err)

	err = sendSearchRequest(node.queryNodeLoopCtx, DIM)
	assert.NoError(t, err)

	time.Sleep(1 * time.Second)

	err = node.Stop()
	assert.NoError(t, err)
}

func TestSearch_SearchMultiSegments(t *testing.T) {
	const N = 10000
	const DIM = 16

	// init queryNode
	collectionID := UniqueID(0)
	segmentID1 := UniqueID(1)
	segmentID2 := UniqueID(2)
	node := newQueryNodeMock()
	initTestMeta(t, node, collectionID, UniqueID(0))

	msFactory, err := newMessageStreamFactory()
	assert.NoError(t, err)

	// start dataSync
	newDS := newDataSyncService(node.queryNodeLoopCtx, node.replica, msFactory, collectionID)
	err = node.addDataSyncService(collectionID, newDS)
	assert.NoError(t, err)
	ds, err := node.getDataSyncService(collectionID)
	assert.NoError(t, err)
	go ds.start()

	// start search service
	node.searchService = newSearchService(node.queryNodeLoopCtx, node.replica, msFactory)
	go node.searchService.start()
	node.searchService.startSearchCollection(collectionID)

	tSafe := node.replica.getTSafe(collectionID)
	assert.NotNil(t, tSafe)
	tSafe.set(1000)

	// load segments
	err = node.replica.addSegment(segmentID1, defaultPartitionID, collectionID, segmentTypeSealed)
	assert.NoError(t, err)
	segment1, err := node.replica.getSegmentByID(segmentID1)
	assert.NoError(t, err)
	err = loadFields(segment1, DIM, N)
	assert.NoError(t, err)

	err = node.replica.addSegment(segmentID2, defaultPartitionID, collectionID, segmentTypeSealed)
	assert.NoError(t, err)
	segment2, err := node.replica.getSegmentByID(segmentID2)
	assert.NoError(t, err)
	err = loadFields(segment2, DIM, N)
	assert.NoError(t, err)

	err = sendSearchRequest(node.queryNodeLoopCtx, DIM)
	assert.NoError(t, err)

	time.Sleep(1 * time.Second)

	err = node.Stop()
	assert.NoError(t, err)
}