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
1897 lines
62 KiB
Go
1897 lines
62 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 proxy
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/require"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
|
"github.com/milvus-io/milvus/internal/proxy/shardclient"
|
|
"github.com/milvus-io/milvus/internal/util/reduce"
|
|
"github.com/milvus-io/milvus/pkg/v3/common"
|
|
"github.com/milvus-io/milvus/pkg/v3/mlog"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/testutils"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
|
|
)
|
|
|
|
func TestQueryTask_all(t *testing.T) {
|
|
var (
|
|
err error
|
|
ctx = context.TODO()
|
|
|
|
qc = NewMixCoordMock()
|
|
qn = getQueryNodeClient()
|
|
|
|
shardsNum = common.DefaultShardsNum
|
|
collectionName = t.Name() + funcutil.GenRandomStr()
|
|
|
|
expr = fmt.Sprintf("%s > 0", testInt64Field)
|
|
hitNum = 10
|
|
)
|
|
|
|
qn.EXPECT().GetComponentStates(mock.Anything, mock.Anything, mock.Anything).Return(nil, nil).Maybe()
|
|
|
|
mgr := shardclient.NewMockShardClientManager(t)
|
|
mgr.EXPECT().GetClient(mock.Anything, mock.Anything).Return(qn, nil).Maybe()
|
|
mgr.EXPECT().GetShardLeaderList(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]string{"mock_qn"}, nil).Maybe()
|
|
mgr.EXPECT().GetShard(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]shardclient.NodeInfo{
|
|
{NodeID: 1, Address: "mock_qn", Serviceable: true},
|
|
}, nil).Maybe()
|
|
mgr.EXPECT().DeprecateShardCache(mock.Anything, mock.Anything).Return().Maybe()
|
|
lb := shardclient.NewLBPolicyImpl(mgr)
|
|
|
|
err = InitMetaCache(ctx, qc)
|
|
assert.NoError(t, err)
|
|
|
|
fieldName2Types := map[string]schemapb.DataType{
|
|
testBoolField: schemapb.DataType_Bool,
|
|
testInt32Field: schemapb.DataType_Int32,
|
|
testInt64Field: schemapb.DataType_Int64,
|
|
testFloatField: schemapb.DataType_Float,
|
|
testDoubleField: schemapb.DataType_Double,
|
|
testFloatVecField: schemapb.DataType_FloatVector,
|
|
}
|
|
if enableMultipleVectorFields {
|
|
fieldName2Types[testBinaryVecField] = schemapb.DataType_BinaryVector
|
|
}
|
|
|
|
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Types, testInt64Field, false)
|
|
marshaledSchema, err := proto.Marshal(schema)
|
|
assert.NoError(t, err)
|
|
|
|
createColT := &createCollectionTask{
|
|
Condition: NewTaskCondition(ctx),
|
|
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
|
|
CollectionName: collectionName,
|
|
Schema: marshaledSchema,
|
|
ShardsNum: shardsNum,
|
|
},
|
|
ctx: ctx,
|
|
mixCoord: qc,
|
|
}
|
|
|
|
require.NoError(t, createColT.OnEnqueue())
|
|
require.NoError(t, createColT.PreExecute(ctx))
|
|
require.NoError(t, createColT.Execute(ctx))
|
|
require.NoError(t, createColT.PostExecute(ctx))
|
|
collectionID, err := globalMetaCache.GetCollectionID(ctx, GetCurDBNameFromContextOrDefault(ctx), collectionName)
|
|
assert.NoError(t, err)
|
|
|
|
status, err := qc.LoadCollection(ctx, &querypb.LoadCollectionRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_LoadCollection,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionID: collectionID,
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
|
|
|
|
t.Run("test query task parameters", func(t *testing.T) {
|
|
task := &queryTask{
|
|
Condition: NewTaskCondition(ctx),
|
|
RetrieveRequest: &internalpb.RetrieveRequest{
|
|
QueryLabel: "query",
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionID: collectionID,
|
|
OutputFieldsId: make([]int64, len(fieldName2Types)),
|
|
},
|
|
ctx: ctx,
|
|
result: &milvuspb.QueryResults{
|
|
Status: merr.Success(),
|
|
FieldsData: []*schemapb.FieldData{},
|
|
},
|
|
request: &milvuspb.QueryRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionName: collectionName,
|
|
Expr: expr,
|
|
QueryParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: IgnoreGrowingKey,
|
|
Value: "false",
|
|
},
|
|
},
|
|
},
|
|
mixCoord: qc,
|
|
lb: lb,
|
|
shardclientMgr: mgr,
|
|
}
|
|
|
|
assert.NoError(t, task.OnEnqueue())
|
|
|
|
// test query task with timeout
|
|
ctx1, cancel1 := context.WithTimeout(ctx, 10*time.Second)
|
|
defer cancel1()
|
|
// before preExecute
|
|
assert.Equal(t, typeutil.ZeroTimestamp, task.TimeoutTimestamp)
|
|
task.ctx = ctx1
|
|
assert.NoError(t, task.PreExecute(ctx))
|
|
|
|
{
|
|
task.mustUsePartitionKey = true
|
|
err := task.PreExecute(ctx)
|
|
assert.Error(t, err)
|
|
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
|
|
task.mustUsePartitionKey = false
|
|
}
|
|
|
|
// after preExecute
|
|
assert.Greater(t, task.TimeoutTimestamp, typeutil.ZeroTimestamp)
|
|
|
|
// check reduce_stop_for_best
|
|
assert.Equal(t, false, task.GetReduceStopForBest())
|
|
task.request.QueryParams = append(task.request.QueryParams, &commonpb.KeyValuePair{
|
|
Key: ReduceStopForBestKey,
|
|
Value: "trxxxx",
|
|
})
|
|
assert.Error(t, task.PreExecute(ctx))
|
|
task.request.QueryParams = task.request.QueryParams[0 : len(task.request.QueryParams)-1]
|
|
|
|
// check parse collection id
|
|
task.request.QueryParams = append(task.request.QueryParams, &commonpb.KeyValuePair{
|
|
Key: CollectionID,
|
|
Value: "trxxxx",
|
|
})
|
|
err := task.PreExecute(ctx)
|
|
assert.Error(t, err)
|
|
task.request.QueryParams = task.request.QueryParams[0 : len(task.request.QueryParams)-1]
|
|
|
|
// check collection id consistency
|
|
task.request.QueryParams = append(task.request.QueryParams, &commonpb.KeyValuePair{
|
|
Key: LimitKey,
|
|
Value: "11",
|
|
})
|
|
task.request.QueryParams = append(task.request.QueryParams, &commonpb.KeyValuePair{
|
|
Key: CollectionID,
|
|
Value: "8080",
|
|
})
|
|
err = task.PreExecute(ctx)
|
|
assert.Error(t, err)
|
|
task.request.QueryParams = make([]*commonpb.KeyValuePair, 0)
|
|
|
|
result1 := &internalpb.RetrieveResults{
|
|
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_RetrieveResult},
|
|
Status: merr.Success(),
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{Data: testutils.GenerateInt64Array(hitNum)},
|
|
},
|
|
},
|
|
}
|
|
|
|
outputFieldIDs := make([]UniqueID, 0, len(fieldName2Types))
|
|
for i := 0; i < len(fieldName2Types); i++ {
|
|
outputFieldIDs = append(outputFieldIDs, int64(common.StartOfUserFieldID+i))
|
|
}
|
|
task.OutputFieldsId = outputFieldIDs
|
|
for fieldName, dataType := range fieldName2Types {
|
|
result1.FieldsData = append(result1.FieldsData, generateFieldData(dataType, fieldName, hitNum))
|
|
}
|
|
result1.FieldsData = append(result1.FieldsData, generateFieldData(schemapb.DataType_Int64, common.TimeStampFieldName, hitNum))
|
|
task.OutputFieldsId = append(task.OutputFieldsId, common.TimeStampField)
|
|
task.ctx = ctx
|
|
qn.ExpectedCalls = nil
|
|
qn.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, nil).Maybe()
|
|
qn.EXPECT().Query(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
|
|
assert.Error(t, task.Execute(ctx))
|
|
|
|
qn.ExpectedCalls = nil
|
|
qn.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, nil).Maybe()
|
|
qn.EXPECT().Query(mock.Anything, mock.Anything).Return(&internalpb.RetrieveResults{
|
|
Status: merr.Status(merr.ErrChannelNotAvailable),
|
|
}, nil)
|
|
err = task.Execute(ctx)
|
|
assert.ErrorIs(t, err, merr.ErrChannelNotAvailable)
|
|
|
|
qn.ExpectedCalls = nil
|
|
qn.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, nil).Maybe()
|
|
qn.EXPECT().Query(mock.Anything, mock.Anything).Return(&internalpb.RetrieveResults{
|
|
Status: &commonpb.Status{
|
|
ErrorCode: commonpb.ErrorCode_UnexpectedError,
|
|
},
|
|
}, nil)
|
|
assert.Error(t, task.Execute(ctx))
|
|
|
|
qn.ExpectedCalls = nil
|
|
qn.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(nil, nil).Maybe()
|
|
qn.EXPECT().Query(mock.Anything, mock.Anything).Return(result1, nil)
|
|
assert.NoError(t, task.Execute(ctx))
|
|
|
|
task.queryParams = &queryParams{
|
|
limit: 100,
|
|
offset: 100,
|
|
}
|
|
assert.NoError(t, task.PostExecute(ctx))
|
|
|
|
for i := 0; i < len(task.result.FieldsData); i++ {
|
|
assert.NotEqual(t, task.result.FieldsData[i].FieldId, common.TimeStampField)
|
|
}
|
|
})
|
|
|
|
t.Run("test query for iterator", func(t *testing.T) {
|
|
qt := &queryTask{
|
|
Condition: NewTaskCondition(ctx),
|
|
RetrieveRequest: &internalpb.RetrieveRequest{
|
|
QueryLabel: "query",
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionID: collectionID,
|
|
OutputFieldsId: make([]int64, len(fieldName2Types)),
|
|
},
|
|
ctx: ctx,
|
|
result: &milvuspb.QueryResults{
|
|
Status: merr.Success(),
|
|
FieldsData: []*schemapb.FieldData{},
|
|
},
|
|
request: &milvuspb.QueryRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionName: collectionName,
|
|
Expr: expr,
|
|
QueryParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: IgnoreGrowingKey,
|
|
Value: "false",
|
|
},
|
|
{
|
|
Key: IteratorField,
|
|
Value: "True",
|
|
},
|
|
},
|
|
},
|
|
mixCoord: qc,
|
|
lb: lb,
|
|
shardclientMgr: mgr,
|
|
resultBuf: &typeutil.ConcurrentSet[*internalpb.RetrieveResults]{},
|
|
}
|
|
// simulate scheduler enqueue task
|
|
enqueTs := uint64(10000)
|
|
qt.SetTs(enqueTs)
|
|
qtErr := qt.PreExecute(context.TODO())
|
|
assert.Nil(t, qtErr)
|
|
assert.True(t, qt.queryParams.isIterator)
|
|
qt.resultBuf.Insert(&internalpb.RetrieveResults{})
|
|
qtErr = qt.PostExecute(context.TODO())
|
|
assert.Nil(t, qtErr)
|
|
// after first page, sessionTs is set
|
|
assert.True(t, qt.result.GetSessionTs() > 0)
|
|
|
|
// next page query task
|
|
qt = &queryTask{
|
|
Condition: NewTaskCondition(ctx),
|
|
RetrieveRequest: &internalpb.RetrieveRequest{
|
|
QueryLabel: "query",
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionID: collectionID,
|
|
OutputFieldsId: make([]int64, len(fieldName2Types)),
|
|
},
|
|
ctx: ctx,
|
|
result: &milvuspb.QueryResults{
|
|
Status: merr.Success(),
|
|
FieldsData: []*schemapb.FieldData{},
|
|
},
|
|
request: &milvuspb.QueryRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionName: collectionName,
|
|
Expr: expr,
|
|
QueryParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: IgnoreGrowingKey,
|
|
Value: "false",
|
|
},
|
|
{
|
|
Key: IteratorField,
|
|
Value: "True",
|
|
},
|
|
},
|
|
GuaranteeTimestamp: enqueTs,
|
|
},
|
|
mixCoord: qc,
|
|
lb: lb,
|
|
shardclientMgr: mgr,
|
|
resultBuf: &typeutil.ConcurrentSet[*internalpb.RetrieveResults]{},
|
|
}
|
|
qtErr = qt.PreExecute(context.TODO())
|
|
assert.Nil(t, qtErr)
|
|
assert.True(t, qt.queryParams.isIterator)
|
|
// from the second page, the mvccTs is set to the sessionTs init in the first page
|
|
assert.Equal(t, enqueTs, qt.GetMvccTimestamp())
|
|
})
|
|
|
|
t.Run("test count(*) with aggregation validation rules", func(t *testing.T) {
|
|
// Test cases for the two validation rule fixes:
|
|
// 1. Aggregation queries exempt from "empty expression needs limit" check (#47326)
|
|
// 2. count(*) pagination only blocked when no GROUP BY (#47326)
|
|
|
|
makeTask := func(exprStr string, outputFields []string, limitValue int64, groupByField string) *queryTask {
|
|
queryParams := []*commonpb.KeyValuePair{
|
|
{Key: IgnoreGrowingKey, Value: "false"},
|
|
}
|
|
if limitValue > 0 {
|
|
queryParams = append(queryParams, &commonpb.KeyValuePair{
|
|
Key: LimitKey,
|
|
Value: fmt.Sprintf("%d", limitValue),
|
|
})
|
|
}
|
|
if groupByField != "" {
|
|
queryParams = append(queryParams, &commonpb.KeyValuePair{
|
|
Key: GroupByFieldsKey,
|
|
Value: groupByField,
|
|
})
|
|
}
|
|
task := &queryTask{
|
|
Condition: NewTaskCondition(ctx),
|
|
RetrieveRequest: &internalpb.RetrieveRequest{
|
|
QueryLabel: "query",
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionID: collectionID,
|
|
OutputFieldsId: make([]int64, 0),
|
|
},
|
|
ctx: ctx,
|
|
result: &milvuspb.QueryResults{
|
|
Status: merr.Success(),
|
|
FieldsData: []*schemapb.FieldData{},
|
|
},
|
|
request: &milvuspb.QueryRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionName: collectionName,
|
|
Expr: exprStr,
|
|
OutputFields: outputFields,
|
|
QueryParams: queryParams,
|
|
},
|
|
mixCoord: qc,
|
|
lb: lb,
|
|
shardclientMgr: mgr,
|
|
}
|
|
task.OnEnqueue()
|
|
return task
|
|
}
|
|
|
|
// Case 1: count(*) + GROUP BY + limit → should PASS
|
|
{
|
|
task := makeTask(
|
|
expr,
|
|
[]string{testInt64Field, "count(*)"},
|
|
100,
|
|
testInt64Field,
|
|
)
|
|
err := task.PreExecute(ctx)
|
|
assert.NoError(t, err, "count(*) with GROUP BY and limit should be allowed")
|
|
}
|
|
|
|
// Case 2: count(*) + GROUP BY + no limit → should PASS
|
|
{
|
|
task := makeTask(
|
|
expr,
|
|
[]string{testInt64Field, "count(*)"},
|
|
0,
|
|
testInt64Field,
|
|
)
|
|
err := task.PreExecute(ctx)
|
|
assert.NoError(t, err, "count(*) with GROUP BY and no limit should be allowed")
|
|
}
|
|
|
|
// Case 3: count(*) + no GROUP BY + limit → should FAIL
|
|
{
|
|
task := makeTask(
|
|
expr,
|
|
[]string{"count(*)"},
|
|
100,
|
|
"",
|
|
)
|
|
err := task.PreExecute(ctx)
|
|
assert.Error(t, err, "count(*) without GROUP BY with limit should fail")
|
|
assert.Contains(t, err.Error(), "count entities with pagination is not allowed")
|
|
}
|
|
|
|
// Case 4: aggregation + empty expr + no limit → should PASS
|
|
{
|
|
task := makeTask(
|
|
"",
|
|
[]string{testInt64Field, "count(*)"},
|
|
0,
|
|
testInt64Field,
|
|
)
|
|
err := task.PreExecute(ctx)
|
|
assert.NoError(t, err, "aggregation with empty expr and no limit should be allowed")
|
|
}
|
|
|
|
// Case 5: non-aggregation + empty expr + no limit → should FAIL
|
|
{
|
|
task := makeTask(
|
|
"",
|
|
[]string{testInt64Field},
|
|
0,
|
|
"",
|
|
)
|
|
err := task.PreExecute(ctx)
|
|
assert.Error(t, err, "non-aggregation with empty expr and no limit should fail")
|
|
assert.Contains(t, err.Error(), "empty expression should be used with limit")
|
|
}
|
|
})
|
|
|
|
t.Run("test order by without limit", func(t *testing.T) {
|
|
// ORDER BY without explicit limit should fail with an error
|
|
task := &queryTask{
|
|
Condition: NewTaskCondition(ctx),
|
|
RetrieveRequest: &internalpb.RetrieveRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionID: collectionID,
|
|
OutputFieldsId: make([]int64, len(fieldName2Types)),
|
|
},
|
|
ctx: ctx,
|
|
request: &milvuspb.QueryRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionName: collectionName,
|
|
Expr: expr,
|
|
QueryParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: IgnoreGrowingKey,
|
|
Value: "false",
|
|
},
|
|
{
|
|
Key: OrderByFieldsKey,
|
|
Value: testInt64Field + ":asc",
|
|
},
|
|
// No limit specified
|
|
},
|
|
},
|
|
mixCoord: qc,
|
|
lb: lb,
|
|
shardclientMgr: mgr,
|
|
}
|
|
assert.NoError(t, task.OnEnqueue())
|
|
err := task.PreExecute(ctx)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "ORDER BY requires explicit limit")
|
|
})
|
|
|
|
t.Run("test order by with limit succeeds", func(t *testing.T) {
|
|
// ORDER BY with explicit limit should pass the validation
|
|
task := &queryTask{
|
|
Condition: NewTaskCondition(ctx),
|
|
RetrieveRequest: &internalpb.RetrieveRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionID: collectionID,
|
|
OutputFieldsId: make([]int64, len(fieldName2Types)),
|
|
},
|
|
ctx: ctx,
|
|
request: &milvuspb.QueryRequest{
|
|
Base: &commonpb.MsgBase{
|
|
MsgType: commonpb.MsgType_Retrieve,
|
|
SourceID: paramtable.GetNodeID(),
|
|
},
|
|
CollectionName: collectionName,
|
|
Expr: expr,
|
|
QueryParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: IgnoreGrowingKey,
|
|
Value: "false",
|
|
},
|
|
{
|
|
Key: OrderByFieldsKey,
|
|
Value: testInt64Field + ":asc",
|
|
},
|
|
{
|
|
Key: LimitKey,
|
|
Value: "10",
|
|
},
|
|
},
|
|
},
|
|
mixCoord: qc,
|
|
lb: lb,
|
|
shardclientMgr: mgr,
|
|
}
|
|
assert.NoError(t, task.OnEnqueue())
|
|
err := task.PreExecute(ctx)
|
|
assert.NoError(t, err)
|
|
})
|
|
}
|
|
|
|
func Test_translateToOutputFieldIDs(t *testing.T) {
|
|
type testCases struct {
|
|
name string
|
|
outputFields []string
|
|
schema *schemapb.CollectionSchema
|
|
expectedError bool
|
|
expectedIDs []int64
|
|
}
|
|
|
|
cases := []testCases{
|
|
{
|
|
name: "empty output fields",
|
|
outputFields: []string{},
|
|
schema: &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.RowIDField,
|
|
Name: common.RowIDFieldName,
|
|
},
|
|
{
|
|
FieldID: 100,
|
|
Name: "ID",
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "Vector",
|
|
},
|
|
},
|
|
},
|
|
expectedError: false,
|
|
expectedIDs: []int64{100},
|
|
},
|
|
{
|
|
name: "nil output fields",
|
|
outputFields: nil,
|
|
schema: &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.RowIDField,
|
|
Name: common.RowIDFieldName,
|
|
},
|
|
{
|
|
FieldID: 100,
|
|
Name: "ID",
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "Vector",
|
|
},
|
|
},
|
|
},
|
|
expectedError: false,
|
|
expectedIDs: []int64{100},
|
|
},
|
|
{
|
|
name: "full list",
|
|
outputFields: []string{"ID", "Vector"},
|
|
schema: &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.RowIDField,
|
|
Name: common.RowIDFieldName,
|
|
},
|
|
{
|
|
FieldID: 100,
|
|
Name: "ID",
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "Vector",
|
|
},
|
|
},
|
|
},
|
|
expectedError: false,
|
|
expectedIDs: []int64{100, 101},
|
|
},
|
|
{
|
|
name: "vector only",
|
|
outputFields: []string{"Vector"},
|
|
schema: &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.RowIDField,
|
|
Name: common.RowIDFieldName,
|
|
},
|
|
{
|
|
FieldID: 100,
|
|
Name: "ID",
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "Vector",
|
|
},
|
|
},
|
|
},
|
|
expectedError: false,
|
|
expectedIDs: []int64{101, 100},
|
|
},
|
|
{
|
|
name: "with field not exist",
|
|
outputFields: []string{"ID", "Vector", "Extra"},
|
|
schema: &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.RowIDField,
|
|
Name: common.RowIDFieldName,
|
|
},
|
|
{
|
|
FieldID: 100,
|
|
Name: "ID",
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "Vector",
|
|
},
|
|
},
|
|
},
|
|
expectedError: true,
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
ids, err := translateToOutputFieldIDs(tc.outputFields, tc.schema)
|
|
if tc.expectedError {
|
|
assert.Error(t, err)
|
|
} else {
|
|
assert.NoError(t, err)
|
|
require.Equal(t, len(tc.expectedIDs), len(ids))
|
|
for idx, expectedID := range tc.expectedIDs {
|
|
assert.Equal(t, expectedID, ids[idx])
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestTaskQuery_functions(t *testing.T) {
|
|
t.Run("test parseQueryParams", func(t *testing.T) {
|
|
tests := []struct {
|
|
description string
|
|
|
|
inKey []string
|
|
inValue []string
|
|
|
|
expectErr bool
|
|
outLimit int64
|
|
outOffset int64
|
|
}{
|
|
{"empty input", []string{}, []string{}, false, typeutil.Unlimited, 0},
|
|
{"valid limit=1", []string{LimitKey}, []string{"1"}, false, 1, 0},
|
|
{"valid limit=1, offset=2", []string{LimitKey, OffsetKey}, []string{"1", "2"}, false, 1, 2},
|
|
{"valid no limit, offset=2", []string{OffsetKey}, []string{"2"}, false, typeutil.Unlimited, 0},
|
|
{"invalid limit str", []string{LimitKey}, []string{"a"}, true, 0, 0},
|
|
{"invalid limit zero", []string{LimitKey}, []string{"0"}, true, 0, 0},
|
|
{"invalid limit negative", []string{LimitKey}, []string{"-1"}, true, 0, 0},
|
|
{"invalid limit 16385", []string{LimitKey}, []string{"16385"}, true, 0, 0},
|
|
{"invalid offset negative", []string{LimitKey, OffsetKey}, []string{"1", "-1"}, true, 0, 0},
|
|
{"invalid offset 16385", []string{LimitKey, OffsetKey}, []string{"1", "16385"}, true, 0, 0},
|
|
{"invalid limit=16384 offset=16384", []string{LimitKey, OffsetKey}, []string{"16384", "16384"}, true, 0, 0},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.description, func(t *testing.T) {
|
|
var inParams []*commonpb.KeyValuePair
|
|
for i := range test.inKey {
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: test.inKey[i],
|
|
Value: test.inValue[i],
|
|
})
|
|
}
|
|
ret, err := parseQueryParams(inParams, false, schemapb.DataType_Int64)
|
|
if test.expectErr {
|
|
assert.Error(t, err)
|
|
assert.Empty(t, ret)
|
|
} else {
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, test.outLimit, ret.limit)
|
|
assert.Equal(t, test.outOffset, ret.offset)
|
|
}
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("test parseQueryParams for reduce type", func(t *testing.T) {
|
|
{
|
|
var inParams []*commonpb.KeyValuePair
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: ReduceStopForBestKey,
|
|
Value: "True",
|
|
})
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: IteratorField,
|
|
Value: "True",
|
|
})
|
|
ret, err := parseQueryParams(inParams, false, schemapb.DataType_Int64)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, reduce.IReduceInOrderForBest, ret.reduceType)
|
|
}
|
|
{
|
|
var inParams []*commonpb.KeyValuePair
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: ReduceStopForBestKey,
|
|
Value: "True",
|
|
})
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: IteratorField,
|
|
Value: "TrueXXXX",
|
|
})
|
|
ret, err := parseQueryParams(inParams, false, schemapb.DataType_Int64)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, ret)
|
|
}
|
|
{
|
|
var inParams []*commonpb.KeyValuePair
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: ReduceStopForBestKey,
|
|
Value: "TrueXXXXX",
|
|
})
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: IteratorField,
|
|
Value: "True",
|
|
})
|
|
ret, err := parseQueryParams(inParams, false, schemapb.DataType_Int64)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, ret)
|
|
}
|
|
{
|
|
var inParams []*commonpb.KeyValuePair
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: ReduceStopForBestKey,
|
|
Value: "True",
|
|
})
|
|
// when not setting iterator tag, ignore reduce_stop_for_best
|
|
ret, err := parseQueryParams(inParams, false, schemapb.DataType_Int64)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, reduce.IReduceNoOrder, ret.reduceType)
|
|
}
|
|
{
|
|
var inParams []*commonpb.KeyValuePair
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: IteratorField,
|
|
Value: "True",
|
|
})
|
|
// when not setting reduce_stop_for_best tag, reduce by keep results in order
|
|
ret, err := parseQueryParams(inParams, false, schemapb.DataType_Int64)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, reduce.IReduceInOrder, ret.reduceType)
|
|
}
|
|
{
|
|
var inParams []*commonpb.KeyValuePair
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: ReduceStopForBestKey,
|
|
Value: "False",
|
|
})
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: IteratorField,
|
|
Value: "True",
|
|
})
|
|
ret, err := parseQueryParams(inParams, false, schemapb.DataType_Int64)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, reduce.IReduceInOrder, ret.reduceType)
|
|
}
|
|
{
|
|
var inParams []*commonpb.KeyValuePair
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: ReduceStopForBestKey,
|
|
Value: "False",
|
|
})
|
|
inParams = append(inParams, &commonpb.KeyValuePair{
|
|
Key: IteratorField,
|
|
Value: "False",
|
|
})
|
|
ret, err := parseQueryParams(inParams, false, schemapb.DataType_Int64)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, reduce.IReduceNoOrder, ret.reduceType)
|
|
}
|
|
})
|
|
|
|
t.Run("test parseQueryParams for query iterator cursor", func(t *testing.T) {
|
|
params := []*commonpb.KeyValuePair{
|
|
{Key: IteratorField, Value: "true"},
|
|
{Key: QueryIterLastPKKey, Value: "7"},
|
|
{Key: QueryIterLastOffsetKey, Value: "2"},
|
|
}
|
|
ret, err := parseQueryParams(params, false, schemapb.DataType_Int64)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, ret.queryIteratorCursor)
|
|
require.NotNil(t, ret.queryIteratorCursor.LastIntPk)
|
|
assert.EqualValues(t, 7, ret.queryIteratorCursor.GetLastIntPk())
|
|
assert.EqualValues(t, 2, ret.queryIteratorCursor.GetLastElementOffset())
|
|
|
|
params = []*commonpb.KeyValuePair{
|
|
{Key: IteratorField, Value: "true"},
|
|
{Key: QueryIterLastPKKey, Value: "pk-7"},
|
|
{Key: QueryIterLastOffsetKey, Value: "2"},
|
|
}
|
|
ret, err = parseQueryParams(params, false, schemapb.DataType_VarChar)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, ret.queryIteratorCursor)
|
|
assert.Equal(t, "pk-7", ret.queryIteratorCursor.GetLastStrPk())
|
|
|
|
params = []*commonpb.KeyValuePair{
|
|
{Key: QueryIterLastPKKey, Value: "7"},
|
|
{Key: QueryIterLastOffsetKey, Value: "2"},
|
|
}
|
|
ret, err = parseQueryParams(params, false, schemapb.DataType_Int64)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, ret)
|
|
|
|
params = []*commonpb.KeyValuePair{
|
|
{Key: IteratorField, Value: "true"},
|
|
{Key: QueryIterLastPKKey, Value: "7"},
|
|
{Key: QueryIterLastOffsetKey, Value: "-1"},
|
|
}
|
|
ret, err = parseQueryParams(params, false, schemapb.DataType_Int64)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, ret)
|
|
})
|
|
|
|
t.Run("test parseQueryIteratorCursor", func(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
params []*commonpb.KeyValuePair
|
|
isIterator bool
|
|
pkDataType schemapb.DataType
|
|
wantNil bool
|
|
wantIntPK int64
|
|
checkIntPK bool
|
|
wantStrPK string
|
|
wantOffset int64
|
|
wantErrSubstr string
|
|
}{
|
|
{
|
|
name: "no cursor params",
|
|
params: nil,
|
|
isIterator: true,
|
|
pkDataType: schemapb.DataType_Int64,
|
|
wantNil: true,
|
|
},
|
|
{
|
|
name: "int64 pk with zero offset",
|
|
params: []*commonpb.KeyValuePair{{Key: QueryIterLastPKKey, Value: "7"}, {Key: QueryIterLastOffsetKey, Value: "0"}},
|
|
isIterator: true,
|
|
pkDataType: schemapb.DataType_Int64,
|
|
wantIntPK: 7,
|
|
checkIntPK: true,
|
|
wantOffset: 0,
|
|
},
|
|
{
|
|
name: "varchar pk",
|
|
params: []*commonpb.KeyValuePair{{Key: QueryIterLastPKKey, Value: "pk-7"}, {Key: QueryIterLastOffsetKey, Value: "2"}},
|
|
isIterator: true,
|
|
pkDataType: schemapb.DataType_VarChar,
|
|
wantStrPK: "pk-7",
|
|
wantOffset: 2,
|
|
},
|
|
{
|
|
name: "cursor params require iterator",
|
|
params: []*commonpb.KeyValuePair{{Key: QueryIterLastPKKey, Value: "7"}, {Key: QueryIterLastOffsetKey, Value: "2"}},
|
|
isIterator: false,
|
|
pkDataType: schemapb.DataType_Int64,
|
|
wantErrSubstr: "invalid query iterator cursor params",
|
|
},
|
|
{
|
|
name: "missing offset",
|
|
params: []*commonpb.KeyValuePair{{Key: QueryIterLastPKKey, Value: "7"}},
|
|
isIterator: true,
|
|
pkDataType: schemapb.DataType_Int64,
|
|
wantErrSubstr: "incomplete query iterator cursor params",
|
|
},
|
|
{
|
|
name: "missing pk",
|
|
params: []*commonpb.KeyValuePair{{Key: QueryIterLastOffsetKey, Value: "2"}},
|
|
isIterator: true,
|
|
pkDataType: schemapb.DataType_Int64,
|
|
wantErrSubstr: "incomplete query iterator cursor params",
|
|
},
|
|
{
|
|
name: "invalid offset",
|
|
params: []*commonpb.KeyValuePair{{Key: QueryIterLastPKKey, Value: "7"}, {Key: QueryIterLastOffsetKey, Value: "abc"}},
|
|
isIterator: true,
|
|
pkDataType: schemapb.DataType_Int64,
|
|
wantErrSubstr: "value for query iterator last element offset is invalid",
|
|
},
|
|
{
|
|
name: "negative offset",
|
|
params: []*commonpb.KeyValuePair{{Key: QueryIterLastPKKey, Value: "7"}, {Key: QueryIterLastOffsetKey, Value: "-1"}},
|
|
isIterator: true,
|
|
pkDataType: schemapb.DataType_Int64,
|
|
wantErrSubstr: "value for query iterator last element offset is invalid",
|
|
},
|
|
{
|
|
name: "invalid int64 pk",
|
|
params: []*commonpb.KeyValuePair{{Key: QueryIterLastPKKey, Value: "abc"}, {Key: QueryIterLastOffsetKey, Value: "2"}},
|
|
isIterator: true,
|
|
pkDataType: schemapb.DataType_Int64,
|
|
wantErrSubstr: "value for query iterator last primary key is invalid",
|
|
},
|
|
{
|
|
name: "unsupported pk type",
|
|
params: []*commonpb.KeyValuePair{{Key: QueryIterLastPKKey, Value: "7"}, {Key: QueryIterLastOffsetKey, Value: "2"}},
|
|
isIterator: true,
|
|
pkDataType: schemapb.DataType_Int32,
|
|
wantErrSubstr: "unsupported primary key type",
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
cursor, err := parseQueryIteratorCursor(test.params, test.isIterator, test.pkDataType)
|
|
if test.wantErrSubstr != "" {
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), test.wantErrSubstr)
|
|
assert.Nil(t, cursor)
|
|
return
|
|
}
|
|
require.NoError(t, err)
|
|
if test.wantNil {
|
|
assert.Nil(t, cursor)
|
|
return
|
|
}
|
|
require.NotNil(t, cursor)
|
|
assert.Equal(t, test.wantOffset, cursor.GetLastElementOffset())
|
|
if test.checkIntPK {
|
|
require.NotNil(t, cursor.LastIntPk)
|
|
assert.Equal(t, test.wantIntPK, cursor.GetLastIntPk())
|
|
}
|
|
if test.wantStrPK != "" {
|
|
require.NotNil(t, cursor.LastStrPk)
|
|
assert.Equal(t, test.wantStrPK, cursor.GetLastStrPk())
|
|
}
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("test reduceRetrieveResults", func(t *testing.T) {
|
|
const (
|
|
Dim = 8
|
|
Int64FieldName = "Int64Field"
|
|
FloatVectorFieldName = "FloatVectorField"
|
|
Int64FieldID = common.StartOfUserFieldID + 1
|
|
FloatVectorFieldID = common.StartOfUserFieldID + 2
|
|
)
|
|
Int64Array := []int64{11, 22}
|
|
FloatVector := []float32{1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 77.0, 88.0}
|
|
|
|
var fieldDataArray1 []*schemapb.FieldData
|
|
fieldDataArray1 = append(fieldDataArray1, getFieldData(Int64FieldName, Int64FieldID, schemapb.DataType_Int64, Int64Array[0:2], 1))
|
|
fieldDataArray1 = append(fieldDataArray1, getFieldData(FloatVectorFieldName, FloatVectorFieldID, schemapb.DataType_FloatVector, FloatVector[0:16], Dim))
|
|
|
|
var fieldDataArray2 []*schemapb.FieldData
|
|
fieldDataArray2 = append(fieldDataArray2, getFieldData(Int64FieldName, Int64FieldID, schemapb.DataType_Int64, Int64Array[0:2], 1))
|
|
fieldDataArray2 = append(fieldDataArray2, getFieldData(FloatVectorFieldName, FloatVectorFieldID, schemapb.DataType_FloatVector, FloatVector[0:16], Dim))
|
|
|
|
t.Run("test nil results", func(t *testing.T) {
|
|
ret, err := reduceRetrieveResults(context.Background(), nil, &queryParams{})
|
|
assert.NoError(t, err)
|
|
assert.Empty(t, ret.GetFieldsData())
|
|
})
|
|
|
|
t.Run("test merge", func(t *testing.T) {
|
|
r1 := &internalpb.RetrieveResults{
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: []int64{1, 3},
|
|
},
|
|
},
|
|
},
|
|
FieldsData: fieldDataArray1,
|
|
}
|
|
r2 := &internalpb.RetrieveResults{
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: []int64{2, 4},
|
|
},
|
|
},
|
|
},
|
|
FieldsData: fieldDataArray2,
|
|
}
|
|
|
|
resultFloat := []float32{
|
|
1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,
|
|
1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,
|
|
11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 77.0, 88.0,
|
|
11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 77.0, 88.0,
|
|
}
|
|
|
|
t.Run("test limited", func(t *testing.T) {
|
|
tests := []struct {
|
|
description string
|
|
limit int64
|
|
}{
|
|
{"limit 1", 1},
|
|
{"limit 2", 2},
|
|
{"limit 3", 3},
|
|
{"limit 4", 4},
|
|
}
|
|
resultField0 := []int64{11, 11, 22, 22}
|
|
for _, test := range tests {
|
|
t.Run(test.description, func(t *testing.T) {
|
|
result, err := reduceRetrieveResults(context.Background(), []*internalpb.RetrieveResults{r1, r2}, &queryParams{limit: test.limit})
|
|
assert.Equal(t, 2, len(result.GetFieldsData()))
|
|
assert.Equal(t, resultField0[0:test.limit], result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
assert.InDeltaSlice(t, resultFloat[0:test.limit*Dim], result.FieldsData[1].GetVectors().GetFloatVector().Data, 10e-10)
|
|
assert.NoError(t, err)
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("test unLimited and maxOutputSize", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QuotaConfig.MaxOutputSize.Key, "1")
|
|
|
|
ids := make([]int64, 100)
|
|
offsets := make([]int64, 100)
|
|
for i := range ids {
|
|
ids[i] = int64(i)
|
|
offsets[i] = int64(i)
|
|
}
|
|
fieldData := getFieldData(Int64FieldName, Int64FieldID, schemapb.DataType_Int64, ids, 1)
|
|
|
|
result := &internalpb.RetrieveResults{
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: ids,
|
|
},
|
|
},
|
|
},
|
|
FieldsData: []*schemapb.FieldData{fieldData},
|
|
}
|
|
|
|
_, err := reduceRetrieveResults(context.Background(), []*internalpb.RetrieveResults{result}, &queryParams{limit: typeutil.Unlimited})
|
|
assert.Error(t, err)
|
|
paramtable.Get().Save(paramtable.Get().QuotaConfig.MaxOutputSize.Key, "1104857600")
|
|
})
|
|
|
|
t.Run("test offset", func(t *testing.T) {
|
|
tests := []struct {
|
|
description string
|
|
offset int64
|
|
}{
|
|
{"offset 0", 0},
|
|
{"offset 1", 1},
|
|
{"offset 2", 2},
|
|
{"offset 3", 3},
|
|
}
|
|
resultField0 := []int64{11, 11, 22, 22}
|
|
for _, test := range tests {
|
|
t.Run(test.description, func(t *testing.T) {
|
|
result, err := reduceRetrieveResults(context.Background(), []*internalpb.RetrieveResults{r1, r2}, &queryParams{limit: 1, offset: test.offset})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 2, len(result.GetFieldsData()))
|
|
assert.Equal(t, resultField0[test.offset:test.offset+1], result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
assert.InDeltaSlice(t, resultFloat[test.offset*Dim:(test.offset+1)*Dim], result.FieldsData[1].GetVectors().GetFloatVector().Data, 10e-10)
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("test stop reduce for best for limit", func(t *testing.T) {
|
|
r1.HasMoreResult = true
|
|
r2.HasMoreResult = false
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: 2, reduceType: reduce.IReduceInOrderForBest})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 2, len(result.GetFieldsData()))
|
|
assert.Equal(t, []int64{11, 11, 22}, result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
len := len(result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
assert.InDeltaSlice(t, resultFloat[0:(len)*Dim], result.FieldsData[1].GetVectors().GetFloatVector().Data, 10e-10)
|
|
})
|
|
|
|
t.Run("test stop reduce for best for limit and offset", func(t *testing.T) {
|
|
r1.HasMoreResult = true
|
|
r2.HasMoreResult = true
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: 1, offset: 1, reduceType: reduce.IReduceInOrderForBest})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 2, len(result.GetFieldsData()))
|
|
assert.Equal(t, []int64{11, 22}, result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
})
|
|
|
|
t.Run("test stop reduce for best for limit and offset", func(t *testing.T) {
|
|
r1.HasMoreResult = false
|
|
r2.HasMoreResult = true
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: 2, offset: 1, reduceType: reduce.IReduceInOrderForBest})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 2, len(result.GetFieldsData()))
|
|
|
|
// we should get 6 result back in total, but only get 4, which means all the result should actually be part of result
|
|
assert.Equal(t, []int64{11, 22, 22}, result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
})
|
|
|
|
t.Run("test stop reduce for best for unlimited set", func(t *testing.T) {
|
|
r1.HasMoreResult = false
|
|
r2.HasMoreResult = false
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: typeutil.Unlimited, reduceType: reduce.IReduceInOrderForBest})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 2, len(result.GetFieldsData()))
|
|
assert.Equal(t, []int64{11, 11, 22, 22}, result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
len := len(result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
assert.InDeltaSlice(t, resultFloat[0:(len)*Dim], result.FieldsData[1].GetVectors().GetFloatVector().Data, 10e-10)
|
|
})
|
|
|
|
t.Run("test stop reduce for best for unlimited set amd offset", func(t *testing.T) {
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: typeutil.Unlimited, offset: 3, reduceType: reduce.IReduceInOrderForBest})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 2, len(result.GetFieldsData()))
|
|
assert.Equal(t, []int64{22}, result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
})
|
|
t.Run("test iterator without setting reduce stop for best", func(t *testing.T) {
|
|
r1.HasMoreResult = true
|
|
r2.HasMoreResult = true
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: 1, reduceType: reduce.IReduceInOrder})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 2, len(result.GetFieldsData()))
|
|
assert.Equal(t, []int64{11}, result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
})
|
|
})
|
|
|
|
t.Run("test SelectMinPK with pk equals MaxInt64 and firstInt true", func(t *testing.T) {
|
|
// Test case to cover pk == minIntPK && firstInt condition in SelectMinPK
|
|
// When the first int64 PK equals math.MaxInt64 (which equals initial minIntPK),
|
|
// and firstInt is true, the condition firstInt || pk < minIntPK evaluates to true
|
|
// This ensures the branch is taken even when pk == minIntPK
|
|
const (
|
|
Dim = 8
|
|
Int64FieldName = "Int64Field"
|
|
FloatVectorFieldName = "FloatVectorField"
|
|
Int64FieldID = common.StartOfUserFieldID + 1
|
|
FloatVectorFieldID = common.StartOfUserFieldID + 2
|
|
)
|
|
maxInt64PK := int64(math.MaxInt64)
|
|
FloatVector := []float32{1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}
|
|
fieldDataMaxPK := getFieldData(Int64FieldName, Int64FieldID, schemapb.DataType_Int64, []int64{maxInt64PK}, 1)
|
|
vectorDataMaxPK := getFieldData(FloatVectorFieldName, FloatVectorFieldID, schemapb.DataType_FloatVector, FloatVector[0:8], Dim)
|
|
|
|
// Create result with MaxInt64 as the first PK to trigger pk == minIntPK && firstInt condition
|
|
// Only the first result has data with pk = maxInt64PK
|
|
rMaxPK := &internalpb.RetrieveResults{
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: []int64{maxInt64PK},
|
|
},
|
|
},
|
|
},
|
|
FieldsData: []*schemapb.FieldData{fieldDataMaxPK, vectorDataMaxPK},
|
|
HasMoreResult: false,
|
|
}
|
|
|
|
// Create a second result with no data (empty result)
|
|
rEmpty := &internalpb.RetrieveResults{
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: []int64{},
|
|
},
|
|
},
|
|
},
|
|
FieldsData: []*schemapb.FieldData{},
|
|
HasMoreResult: false,
|
|
}
|
|
|
|
// Test with MaxInt64 first to ensure it's handled correctly when pk == minIntPK && firstInt
|
|
// The reduced result should include the maxInt64PK result
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{rMaxPK, rEmpty},
|
|
&queryParams{limit: typeutil.Unlimited})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 2, len(result.GetFieldsData()))
|
|
// Should include the maxInt64PK result
|
|
assert.Equal(t, []int64{maxInt64PK}, result.GetFieldsData()[0].GetScalars().GetLongData().Data)
|
|
})
|
|
|
|
t.Run("test element-level reduce", func(t *testing.T) {
|
|
// Helper to create element-level retrieve result
|
|
makeElementLevelResult := func(pks []int64, elemIndices [][]int32) *internalpb.RetrieveResults {
|
|
elemIdxList := make([]*internalpb.ElementIndices, len(elemIndices))
|
|
for i, indices := range elemIndices {
|
|
elemIdxList[i] = &internalpb.ElementIndices{Indices: indices}
|
|
}
|
|
return &internalpb.RetrieveResults{
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{Data: pks},
|
|
},
|
|
},
|
|
FieldsData: []*schemapb.FieldData{
|
|
getFieldData(Int64FieldName, Int64FieldID, schemapb.DataType_Int64, pks, 1),
|
|
},
|
|
ElementLevel: true,
|
|
ElementIndices: elemIdxList,
|
|
}
|
|
}
|
|
|
|
t.Run("basic element-level merge", func(t *testing.T) {
|
|
r1 := makeElementLevelResult([]int64{1, 3}, [][]int32{{0, 1}, {2}})
|
|
r2 := makeElementLevelResult([]int64{2, 4}, [][]int32{{0}, {1, 2}})
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: typeutil.Unlimited})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 4, len(result.GetElementIndices()))
|
|
// Verify element_indices are converted to int64
|
|
assert.Equal(t, []int64{0, 1}, result.GetElementIndices()[0].GetIndices().GetData())
|
|
assert.Equal(t, []int64{0}, result.GetElementIndices()[1].GetIndices().GetData())
|
|
assert.Equal(t, []int64{2}, result.GetElementIndices()[2].GetIndices().GetData())
|
|
assert.Equal(t, []int64{1, 2}, result.GetElementIndices()[3].GetIndices().GetData())
|
|
})
|
|
|
|
t.Run("element-level limit counts elements not docs", func(t *testing.T) {
|
|
// doc1 has 3 elements, doc2 has 2 elements
|
|
r := makeElementLevelResult([]int64{1, 2}, [][]int32{{0, 1, 2}, {0, 1}})
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r},
|
|
&queryParams{limit: 3}) // limit=3 elements
|
|
assert.NoError(t, err)
|
|
// Should get doc1 (3 elements) only
|
|
assert.Equal(t, 1, len(result.GetFieldsData()[0].GetScalars().GetLongData().GetData()))
|
|
assert.Equal(t, int64(1), result.GetFieldsData()[0].GetScalars().GetLongData().GetData()[0])
|
|
assert.Equal(t, 1, len(result.GetElementIndices()))
|
|
})
|
|
|
|
t.Run("element-level limit allows partial fill", func(t *testing.T) {
|
|
// doc1 has 2 elements, doc2 has 2 elements
|
|
r := makeElementLevelResult([]int64{1, 2}, [][]int32{{0, 1}, {0, 1}})
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r},
|
|
&queryParams{limit: 3}) // limit=3 elements
|
|
assert.NoError(t, err)
|
|
// Behavior: process docs while availableCount < limit
|
|
// doc1 (2 elements): 0 < 3, process, count=2
|
|
// doc2 (2 elements): 2 < 3, process, count=4
|
|
// Both docs are included because the check is done before processing each doc
|
|
assert.Equal(t, 2, len(result.GetFieldsData()[0].GetScalars().GetLongData().GetData()))
|
|
})
|
|
|
|
t.Run("element-level with offset", func(t *testing.T) {
|
|
r := makeElementLevelResult([]int64{1, 2, 3}, [][]int32{{0}, {1}, {2}})
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r},
|
|
&queryParams{limit: 2, offset: 1}) // skip 1 doc, get 2 docs
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 2, len(result.GetFieldsData()[0].GetScalars().GetLongData().GetData()))
|
|
assert.Equal(t, []int64{2, 3}, result.GetFieldsData()[0].GetScalars().GetLongData().GetData())
|
|
})
|
|
|
|
t.Run("element-level offset should skip by element count not doc count", func(t *testing.T) {
|
|
// Doc layout:
|
|
// pk=1: 3 elements [0,1,2]
|
|
// pk=2: 2 elements [0,1]
|
|
// pk=3: 1 element [0]
|
|
// Total: 6 elements
|
|
//
|
|
// Query: offset=3, limit=3
|
|
// Correct (element-level offset):
|
|
// Skip 3 elements (all from pk=1) → return pk=2 (2 elem) + pk=3 (1 elem) = 3 elements
|
|
// Bug (doc-level offset):
|
|
// Skip 3 docs (pk=1, pk=2, pk=3) → nothing left → empty result
|
|
r := makeElementLevelResult([]int64{1, 2, 3}, [][]int32{{0, 1, 2}, {0, 1}, {0}})
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r},
|
|
&queryParams{limit: 3, offset: 3})
|
|
assert.NoError(t, err)
|
|
pks := result.GetFieldsData()[0].GetScalars().GetLongData().GetData()
|
|
assert.Equal(t, []int64{2, 3}, pks, "offset=3 should skip 3 elements (all of pk=1), returning pk=2 and pk=3")
|
|
assert.Equal(t, 2, len(result.GetElementIndices()))
|
|
assert.Equal(t, []int64{0, 1}, result.GetElementIndices()[0].GetIndices().GetData())
|
|
assert.Equal(t, []int64{0}, result.GetElementIndices()[1].GetIndices().GetData())
|
|
})
|
|
|
|
t.Run("element-level offset falls mid-document trims elements", func(t *testing.T) {
|
|
// Doc layout:
|
|
// pk=1: 3 elements [0,1,2]
|
|
// pk=2: 2 elements [3,4]
|
|
// Total: 5 elements
|
|
//
|
|
// Query: offset=2, limit=10
|
|
// offset=2 falls in the middle of pk=1 (which has 3 elements).
|
|
// Should skip elements [0,1] of pk=1, keep element [2].
|
|
// Result: pk=1 with trimmed indices [2], pk=2 with [3,4] → 3 elements total
|
|
r := makeElementLevelResult([]int64{1, 2}, [][]int32{{0, 1, 2}, {3, 4}})
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r},
|
|
&queryParams{limit: 10, offset: 2})
|
|
assert.NoError(t, err)
|
|
pks := result.GetFieldsData()[0].GetScalars().GetLongData().GetData()
|
|
assert.Equal(t, []int64{1, 2}, pks, "both docs should be present since offset lands mid-doc on pk=1")
|
|
assert.Equal(t, 2, len(result.GetElementIndices()))
|
|
// pk=1: original [0,1,2], trimmed first 2 → [2]
|
|
assert.Equal(t, []int64{2}, result.GetElementIndices()[0].GetIndices().GetData())
|
|
// pk=2: unchanged [3,4]
|
|
assert.Equal(t, []int64{3, 4}, result.GetElementIndices()[1].GetIndices().GetData())
|
|
})
|
|
|
|
t.Run("element-level offset+limit pagination consistency", func(t *testing.T) {
|
|
// Verify that two consecutive pages cover all elements without gap or overlap.
|
|
// Doc layout (across 2 shards):
|
|
// shard0: pk=1 (2 elements [0,1]), pk=3 (1 element [0])
|
|
// shard1: pk=2 (3 elements [0,1,2]), pk=4 (2 elements [0,1])
|
|
// Merged by PK: pk=1(2), pk=2(3), pk=3(1), pk=4(2) → total 8 elements
|
|
r1 := makeElementLevelResult([]int64{1, 3}, [][]int32{{0, 1}, {0}})
|
|
r2 := makeElementLevelResult([]int64{2, 4}, [][]int32{{0, 1, 2}, {0, 1}})
|
|
|
|
// Page 1: offset=0, limit=5 → first 5 elements
|
|
page1, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: 5, offset: 0})
|
|
assert.NoError(t, err)
|
|
page1PKs := page1.GetFieldsData()[0].GetScalars().GetLongData().GetData()
|
|
|
|
// Page 2: offset=5, limit=5 → next 3 elements (only 8 total)
|
|
// Recreate results since cursors are consumed
|
|
r1 = makeElementLevelResult([]int64{1, 3}, [][]int32{{0, 1}, {0}})
|
|
r2 = makeElementLevelResult([]int64{2, 4}, [][]int32{{0, 1, 2}, {0, 1}})
|
|
page2, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: 5, offset: 5})
|
|
assert.NoError(t, err)
|
|
page2PKs := page2.GetFieldsData()[0].GetScalars().GetLongData().GetData()
|
|
|
|
// Count total elements across both pages
|
|
page1ElemCount := 0
|
|
for _, ei := range page1.GetElementIndices() {
|
|
page1ElemCount += len(ei.GetIndices().GetData())
|
|
}
|
|
page2ElemCount := 0
|
|
for _, ei := range page2.GetElementIndices() {
|
|
page2ElemCount += len(ei.GetIndices().GetData())
|
|
}
|
|
|
|
// The two pages together should cover exactly 8 elements total
|
|
assert.Equal(t, 8, page1ElemCount+page2ElemCount,
|
|
"page1 (%v, %d elems) + page2 (%v, %d elems) should cover all 8 elements",
|
|
page1PKs, page1ElemCount, page2PKs, page2ElemCount)
|
|
})
|
|
|
|
t.Run("element-level validation: inconsistent flag", func(t *testing.T) {
|
|
r1 := makeElementLevelResult([]int64{1}, [][]int32{{0, 1}})
|
|
r2 := &internalpb.RetrieveResults{
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{Data: []int64{2}},
|
|
},
|
|
},
|
|
FieldsData: []*schemapb.FieldData{
|
|
getFieldData(Int64FieldName, Int64FieldID, schemapb.DataType_Int64, []int64{2}, 1),
|
|
},
|
|
ElementLevel: false, // inconsistent
|
|
}
|
|
_, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: typeutil.Unlimited})
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "inconsistent element-level flag")
|
|
})
|
|
|
|
t.Run("element-level validation: mismatched element_indices length", func(t *testing.T) {
|
|
r := &internalpb.RetrieveResults{
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{Data: []int64{1, 2}}, // 2 ids
|
|
},
|
|
},
|
|
FieldsData: []*schemapb.FieldData{
|
|
getFieldData(Int64FieldName, Int64FieldID, schemapb.DataType_Int64, []int64{1, 2}, 1),
|
|
},
|
|
ElementLevel: true,
|
|
ElementIndices: []*internalpb.ElementIndices{
|
|
{Indices: []int32{0}}, // only 1 element_indices, should be 2
|
|
},
|
|
}
|
|
_, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r},
|
|
&queryParams{limit: typeutil.Unlimited})
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "element_indices length")
|
|
})
|
|
|
|
t.Run("element-level type conversion", func(t *testing.T) {
|
|
r := makeElementLevelResult([]int64{1}, [][]int32{{0, 100, 200}})
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r},
|
|
&queryParams{limit: typeutil.Unlimited})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 1, len(result.GetElementIndices()))
|
|
// Verify int32 -> int64 conversion
|
|
indices := result.GetElementIndices()[0].GetIndices().GetData()
|
|
assert.Equal(t, []int64{0, 100, 200}, indices)
|
|
})
|
|
|
|
t.Run("element-level with IReduceInOrderForBest", func(t *testing.T) {
|
|
// 2 results from different segments, both have HasMoreResult=true
|
|
// r1: PK=1 (2 elems), PK=3 (1 elem)
|
|
// r2: PK=2 (3 elems), PK=4 (1 elem)
|
|
// Merged order by PK: 1,2,3,4
|
|
// IReduceInOrderForBest does not use input limit for elements
|
|
// Stops when r1 drains after PK3 (ShouldStopWhenDrained=true)
|
|
r1 := makeElementLevelResult([]int64{1, 3}, [][]int32{{0, 1}, {2}})
|
|
r1.HasMoreResult = true
|
|
r2 := makeElementLevelResult([]int64{2, 4}, [][]int32{{0, 1, 2}, {1}})
|
|
r2.HasMoreResult = true
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: 4, reduceType: reduce.IReduceInOrderForBest})
|
|
assert.NoError(t, err)
|
|
// IReduceInOrderForBest does not enforce element limit (ShouldUseInputLimit=false)
|
|
// It stops when one result is drained (ShouldStopWhenDrained=true)
|
|
// r1 drains after PK3, so result is {1, 2, 3}
|
|
pks := result.GetFieldsData()[0].GetScalars().GetLongData().GetData()
|
|
assert.Equal(t, []int64{1, 2, 3}, pks)
|
|
assert.Equal(t, 3, len(result.GetElementIndices()))
|
|
})
|
|
|
|
t.Run("element-level with IReduceInOrderForBest drains one result", func(t *testing.T) {
|
|
// r1 has no more (HasMoreResult=false), r2 has more (HasMoreResult=true)
|
|
// drainResult only triggers when HasMoreResult=true, so r1 draining
|
|
// does not stop the loop — it's safe to continue with r2's remaining PKs
|
|
r1 := makeElementLevelResult([]int64{1}, [][]int32{{0, 1}})
|
|
r1.HasMoreResult = false
|
|
r2 := makeElementLevelResult([]int64{2, 4}, [][]int32{{0}, {1, 2}})
|
|
r2.HasMoreResult = true
|
|
result, err := reduceRetrieveResults(context.Background(),
|
|
[]*internalpb.RetrieveResults{r1, r2},
|
|
&queryParams{limit: 10, reduceType: reduce.IReduceInOrderForBest})
|
|
assert.NoError(t, err)
|
|
// r1 HasMoreResult=false, so drain does not trigger stop
|
|
// All PKs are returned: 1, 2, 4
|
|
pks := result.GetFieldsData()[0].GetScalars().GetLongData().GetData()
|
|
assert.Equal(t, []int64{1, 2, 4}, pks)
|
|
assert.Equal(t, 3, len(result.GetElementIndices()))
|
|
})
|
|
})
|
|
})
|
|
}
|
|
|
|
func getFieldData(fieldName string, fieldID int64, fieldType schemapb.DataType, fieldValue interface{}, dim int64) *schemapb.FieldData {
|
|
var fieldData *schemapb.FieldData
|
|
switch fieldType {
|
|
case schemapb.DataType_Bool:
|
|
fieldData = &schemapb.FieldData{
|
|
Type: schemapb.DataType_Bool,
|
|
FieldName: fieldName,
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_BoolData{
|
|
BoolData: &schemapb.BoolArray{
|
|
Data: fieldValue.([]bool),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
FieldId: fieldID,
|
|
}
|
|
case schemapb.DataType_Int32:
|
|
fieldData = &schemapb.FieldData{
|
|
Type: schemapb.DataType_Int32,
|
|
FieldName: fieldName,
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_IntData{
|
|
IntData: &schemapb.IntArray{
|
|
Data: fieldValue.([]int32),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
FieldId: fieldID,
|
|
}
|
|
case schemapb.DataType_Int64:
|
|
fieldData = &schemapb.FieldData{
|
|
Type: schemapb.DataType_Int64,
|
|
FieldName: fieldName,
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_LongData{
|
|
LongData: &schemapb.LongArray{
|
|
Data: fieldValue.([]int64),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
FieldId: fieldID,
|
|
}
|
|
case schemapb.DataType_Float:
|
|
fieldData = &schemapb.FieldData{
|
|
Type: schemapb.DataType_Float,
|
|
FieldName: fieldName,
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_FloatData{
|
|
FloatData: &schemapb.FloatArray{
|
|
Data: fieldValue.([]float32),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
FieldId: fieldID,
|
|
}
|
|
case schemapb.DataType_Double:
|
|
fieldData = &schemapb.FieldData{
|
|
Type: schemapb.DataType_Double,
|
|
FieldName: fieldName,
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_DoubleData{
|
|
DoubleData: &schemapb.DoubleArray{
|
|
Data: fieldValue.([]float64),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
FieldId: fieldID,
|
|
}
|
|
case schemapb.DataType_VarChar:
|
|
fieldData = &schemapb.FieldData{
|
|
Type: schemapb.DataType_VarChar,
|
|
FieldName: fieldName,
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_StringData{
|
|
StringData: &schemapb.StringArray{
|
|
Data: fieldValue.([]string),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
FieldId: fieldID,
|
|
}
|
|
|
|
case schemapb.DataType_BinaryVector:
|
|
fieldData = &schemapb.FieldData{
|
|
Type: schemapb.DataType_BinaryVector,
|
|
FieldName: fieldName,
|
|
Field: &schemapb.FieldData_Vectors{
|
|
Vectors: &schemapb.VectorField{
|
|
Dim: dim,
|
|
Data: &schemapb.VectorField_BinaryVector{
|
|
BinaryVector: fieldValue.([]byte),
|
|
},
|
|
},
|
|
},
|
|
FieldId: fieldID,
|
|
}
|
|
case schemapb.DataType_FloatVector:
|
|
fieldData = &schemapb.FieldData{
|
|
Type: schemapb.DataType_FloatVector,
|
|
FieldName: fieldName,
|
|
Field: &schemapb.FieldData_Vectors{
|
|
Vectors: &schemapb.VectorField{
|
|
Dim: dim,
|
|
Data: &schemapb.VectorField_FloatVector{
|
|
FloatVector: &schemapb.FloatArray{
|
|
Data: fieldValue.([]float32),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
FieldId: fieldID,
|
|
}
|
|
default:
|
|
mlog.Warn(context.TODO(), "not supported field type", mlog.String("fieldType", fieldType.String()))
|
|
}
|
|
|
|
return fieldData
|
|
}
|
|
|
|
func Test_filterSystemFields(t *testing.T) {
|
|
outputFieldIDs := []UniqueID{common.RowIDField, common.TimeStampField, common.StartOfUserFieldID}
|
|
filtered := filterSystemFields(outputFieldIDs)
|
|
assert.ElementsMatch(t, []UniqueID{common.StartOfUserFieldID}, filtered)
|
|
}
|
|
|
|
func Test_createCntPlan(t *testing.T) {
|
|
t.Run("plan without filter", func(t *testing.T) {
|
|
plan, err := createCntPlan("", nil, nil)
|
|
assert.NoError(t, err)
|
|
assert.True(t, plan.GetQuery().GetIsCount())
|
|
assert.Nil(t, plan.GetQuery().GetPredicates())
|
|
})
|
|
|
|
t.Run("invalid schema", func(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "a",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
},
|
|
}
|
|
schemaHelper, err := typeutil.CreateSchemaHelper(schema)
|
|
require.NoError(t, err)
|
|
plan, err := createCntPlan("a > 4", schemaHelper, nil)
|
|
assert.NoError(t, err)
|
|
assert.True(t, plan.GetQuery().GetIsCount())
|
|
assert.NotNil(t, plan.GetQuery().GetPredicates())
|
|
})
|
|
}
|
|
|
|
func Test_queryTask_createPlan(t *testing.T) {
|
|
collSchema := newTestSchema()
|
|
t.Run("query without expression", func(t *testing.T) {
|
|
schema := newSchemaInfo(collSchema)
|
|
tsk := &queryTask{
|
|
request: &milvuspb.QueryRequest{
|
|
OutputFields: []string{"Int64"},
|
|
},
|
|
schema: schema,
|
|
}
|
|
err := tsk.createPlan(context.TODO())
|
|
assert.Error(t, err)
|
|
})
|
|
|
|
t.Run("invalid expression", func(t *testing.T) {
|
|
schema := newSchemaInfo(collSchema)
|
|
|
|
tsk := &queryTask{
|
|
schema: schema,
|
|
request: &milvuspb.QueryRequest{
|
|
OutputFields: []string{"a"},
|
|
Expr: "b > 2",
|
|
},
|
|
}
|
|
err := tsk.createPlan(context.TODO())
|
|
assert.Error(t, err)
|
|
})
|
|
|
|
t.Run("invalid output fields", func(t *testing.T) {
|
|
schema := newSchemaInfo(collSchema)
|
|
|
|
tsk := &queryTask{
|
|
schema: schema,
|
|
request: &milvuspb.QueryRequest{
|
|
OutputFields: []string{"b"},
|
|
Expr: "a > 2",
|
|
},
|
|
}
|
|
err := tsk.createPlan(context.TODO())
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
func TestQueryTask_IDs2Expr(t *testing.T) {
|
|
fieldName := "pk"
|
|
intIDs := &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: []int64{1, 2, 3, 4, 5},
|
|
},
|
|
},
|
|
}
|
|
stringIDs := &schemapb.IDs{
|
|
IdField: &schemapb.IDs_StrId{
|
|
StrId: &schemapb.StringArray{
|
|
Data: []string{"a", "b", "c"},
|
|
},
|
|
},
|
|
}
|
|
idExpr := IDs2Expr(fieldName, intIDs)
|
|
expectIDExpr := "pk in [ 1, 2, 3, 4, 5 ]"
|
|
assert.Equal(t, expectIDExpr, idExpr)
|
|
|
|
strExpr := IDs2Expr(fieldName, stringIDs)
|
|
expectStrExpr := "pk in [ \"a\", \"b\", \"c\" ]"
|
|
assert.Equal(t, expectStrExpr, strExpr)
|
|
}
|
|
|
|
func TestQueryTask_CanSkipAllocTimestamp(t *testing.T) {
|
|
dbName := "test_query"
|
|
collName := "test_skip_alloc_timestamp"
|
|
collID := UniqueID(111)
|
|
mockMetaCache := NewMockCache(t)
|
|
globalMetaCache = mockMetaCache
|
|
|
|
t.Run("default consistency level", func(t *testing.T) {
|
|
qt := &queryTask{
|
|
request: &milvuspb.QueryRequest{
|
|
Base: nil,
|
|
DbName: dbName,
|
|
CollectionName: collName,
|
|
UseDefaultConsistency: true,
|
|
},
|
|
}
|
|
mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, nil)
|
|
mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
|
|
&collectionInfo{
|
|
collID: collID,
|
|
consistencyLevel: commonpb.ConsistencyLevel_Eventually,
|
|
}, nil).Once()
|
|
|
|
skip := qt.CanSkipAllocTimestamp()
|
|
assert.True(t, skip)
|
|
|
|
mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
|
|
&collectionInfo{
|
|
collID: collID,
|
|
consistencyLevel: commonpb.ConsistencyLevel_Bounded,
|
|
}, nil).Once()
|
|
skip = qt.CanSkipAllocTimestamp()
|
|
assert.True(t, skip)
|
|
|
|
mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
|
|
&collectionInfo{
|
|
collID: collID,
|
|
consistencyLevel: commonpb.ConsistencyLevel_Strong,
|
|
}, nil).Once()
|
|
skip = qt.CanSkipAllocTimestamp()
|
|
assert.False(t, skip)
|
|
})
|
|
|
|
t.Run("request consistency level", func(t *testing.T) {
|
|
mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
|
|
&collectionInfo{
|
|
collID: collID,
|
|
consistencyLevel: commonpb.ConsistencyLevel_Eventually,
|
|
}, nil).Times(3)
|
|
|
|
qt := &queryTask{
|
|
request: &milvuspb.QueryRequest{
|
|
Base: nil,
|
|
DbName: dbName,
|
|
CollectionName: collName,
|
|
UseDefaultConsistency: false,
|
|
ConsistencyLevel: commonpb.ConsistencyLevel_Eventually,
|
|
},
|
|
}
|
|
|
|
skip := qt.CanSkipAllocTimestamp()
|
|
assert.True(t, skip)
|
|
|
|
qt.request.ConsistencyLevel = commonpb.ConsistencyLevel_Bounded
|
|
skip = qt.CanSkipAllocTimestamp()
|
|
assert.True(t, skip)
|
|
|
|
qt.request.ConsistencyLevel = commonpb.ConsistencyLevel_Strong
|
|
skip = qt.CanSkipAllocTimestamp()
|
|
assert.False(t, skip)
|
|
})
|
|
|
|
t.Run("legacy_guarantee_ts", func(t *testing.T) {
|
|
qt := &queryTask{
|
|
request: &milvuspb.QueryRequest{
|
|
Base: nil,
|
|
DbName: dbName,
|
|
CollectionName: collName,
|
|
UseDefaultConsistency: false,
|
|
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
|
|
},
|
|
}
|
|
|
|
skip := qt.CanSkipAllocTimestamp()
|
|
assert.False(t, skip)
|
|
|
|
qt.request.GuaranteeTimestamp = 1 // eventually
|
|
skip = qt.CanSkipAllocTimestamp()
|
|
assert.True(t, skip)
|
|
|
|
qt.request.GuaranteeTimestamp = 2 // bounded
|
|
skip = qt.CanSkipAllocTimestamp()
|
|
assert.True(t, skip)
|
|
})
|
|
|
|
t.Run("failed", func(t *testing.T) {
|
|
mockMetaCache.ExpectedCalls = nil
|
|
mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, nil)
|
|
mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
|
|
nil, errors.New("mock error")).Once()
|
|
|
|
qt := &queryTask{
|
|
request: &milvuspb.QueryRequest{
|
|
Base: nil,
|
|
DbName: dbName,
|
|
CollectionName: collName,
|
|
UseDefaultConsistency: true,
|
|
ConsistencyLevel: commonpb.ConsistencyLevel_Eventually,
|
|
},
|
|
}
|
|
|
|
skip := qt.CanSkipAllocTimestamp()
|
|
assert.False(t, skip)
|
|
|
|
mockMetaCache.ExpectedCalls = nil
|
|
mockMetaCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(collID, errors.New("mock error"))
|
|
mockMetaCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(
|
|
&collectionInfo{
|
|
collID: collID,
|
|
consistencyLevel: commonpb.ConsistencyLevel_Eventually,
|
|
}, nil)
|
|
|
|
skip = qt.CanSkipAllocTimestamp()
|
|
assert.False(t, skip)
|
|
|
|
qt2 := &queryTask{
|
|
request: &milvuspb.QueryRequest{
|
|
Base: nil,
|
|
DbName: dbName,
|
|
CollectionName: collName,
|
|
UseDefaultConsistency: false,
|
|
ConsistencyLevel: commonpb.ConsistencyLevel_Eventually,
|
|
},
|
|
}
|
|
|
|
skip = qt2.CanSkipAllocTimestamp()
|
|
assert.True(t, skip)
|
|
})
|
|
}
|