Files
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
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

9001 lines
281 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 (
"bytes"
"context"
"encoding/binary"
"fmt"
"math/rand"
"strconv"
"strings"
"testing"
"time"
"github.com/bytedance/mockey"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"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/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/allocator"
"github.com/milvus-io/milvus/internal/json"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proxy/shardclient"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/util/function/embedding"
"github.com/milvus-io/milvus/internal/util/indexparamcheck"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
"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/proto/rootcoordpb"
"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
"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/metric"
"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/timerecord"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
"github.com/milvus-io/milvus/pkg/v3/util/uniquegenerator"
)
// TODO(dragondriver): add more test cases
const (
maxTestStringLen = 100
testBoolField = "bool"
testInt32Field = "int32"
testInt64Field = "int64"
testFloatField = "float"
testDoubleField = "double"
testVarCharField = "varChar"
testFloatVecField = "fvec"
testBinaryVecField = "bvec"
testFloat16VecField = "f16vec"
testBFloat16VecField = "bf16vec"
testStructArrayField = "structArray"
testTTLField = "ttl"
testGeometryField = "geometry"
testVecDim = 128
testMaxVarCharLength = 100
)
func genCollectionSchema(collectionName string) *schemapb.CollectionSchema {
return constructCollectionSchemaWithAllType(
testBoolField,
testInt32Field,
testInt64Field,
testFloatField,
testDoubleField,
testFloatVecField,
testBinaryVecField,
testFloat16VecField,
testBFloat16VecField,
testStructArrayField,
testGeometryField,
testVecDim,
collectionName)
}
func constructCollectionSchema(
int64Field, floatVecField string,
dim int,
collectionName string,
) *schemapb.CollectionSchema {
pk := &schemapb.FieldSchema{
FieldID: 100,
Name: int64Field,
IsPrimaryKey: true,
Description: "",
DataType: schemapb.DataType_Int64,
TypeParams: nil,
IndexParams: nil,
AutoID: true,
}
fVec := &schemapb.FieldSchema{
FieldID: 101,
Name: floatVecField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(dim),
},
},
IndexParams: nil,
AutoID: false,
}
return &schemapb.CollectionSchema{
Name: collectionName,
Description: "",
AutoID: false,
Fields: []*schemapb.FieldSchema{
pk,
fVec,
},
Properties: []*commonpb.KeyValuePair{
{
Key: common.CollectionTTLConfigKey,
Value: "15",
},
},
}
}
func constructCollectionSchemaEnableDynamicSchema(
int64Field, floatVecField string,
dim int,
collectionName string,
) *schemapb.CollectionSchema {
pk := &schemapb.FieldSchema{
FieldID: 0,
Name: int64Field,
IsPrimaryKey: true,
Description: "",
DataType: schemapb.DataType_Int64,
TypeParams: nil,
IndexParams: nil,
AutoID: true,
}
fVec := &schemapb.FieldSchema{
FieldID: 0,
Name: floatVecField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(dim),
},
},
IndexParams: nil,
AutoID: false,
}
return &schemapb.CollectionSchema{
Name: collectionName,
Description: "",
AutoID: false,
EnableDynamicField: true,
Fields: []*schemapb.FieldSchema{
pk,
fVec,
},
}
}
func ConstructCollectionSchemaWithPartitionKey(collectionName string, fieldName2DataType map[string]schemapb.DataType, primaryFieldName string, partitionKeyFieldName string, autoID bool) *schemapb.CollectionSchema {
schema := constructCollectionSchemaByDataType(collectionName, fieldName2DataType, primaryFieldName, autoID)
for _, field := range schema.Fields {
if field.Name == partitionKeyFieldName {
field.IsPartitionKey = true
}
}
return schema
}
func constructCollectionSchemaByDataType(collectionName string, fieldName2DataType map[string]schemapb.DataType, primaryFieldName string, autoID bool) *schemapb.CollectionSchema {
fieldsSchema := make([]*schemapb.FieldSchema, 0)
idx := int64(100)
for fieldName, dataType := range fieldName2DataType {
fieldSchema := &schemapb.FieldSchema{
FieldID: idx,
Name: fieldName,
DataType: dataType,
}
idx++
if typeutil.IsVectorType(dataType) {
fieldSchema.TypeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(testVecDim),
},
}
}
if dataType == schemapb.DataType_VarChar {
fieldSchema.TypeParams = []*commonpb.KeyValuePair{
{
Key: common.MaxLengthKey,
Value: strconv.Itoa(testMaxVarCharLength),
},
}
}
if fieldName == primaryFieldName {
fieldSchema.IsPrimaryKey = true
fieldSchema.AutoID = autoID
}
fieldsSchema = append(fieldsSchema, fieldSchema)
}
return &schemapb.CollectionSchema{
Name: collectionName,
Fields: fieldsSchema,
}
}
func newTextSchemaForStorageV3Test(collectionName string) *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Name: collectionName,
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: testInt64Field, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_Text},
{
FieldID: 102,
Name: testFloatVecField,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: strconv.Itoa(testVecDim)},
},
},
},
}
}
func constructCollectionSchemaWithTTLField(collectionName string, ttlField string) *schemapb.CollectionSchema {
pk := &schemapb.FieldSchema{
FieldID: 100,
Name: int64Field,
IsPrimaryKey: true,
Description: "",
DataType: schemapb.DataType_Int64,
TypeParams: nil,
IndexParams: nil,
AutoID: true,
}
tz := &schemapb.FieldSchema{
FieldID: 101,
Name: ttlField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_Timestamptz,
TypeParams: nil,
IndexParams: nil,
AutoID: false,
}
fVec := &schemapb.FieldSchema{
FieldID: 102,
Name: floatVecField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(dim),
},
},
IndexParams: nil,
AutoID: false,
}
return &schemapb.CollectionSchema{
Name: collectionName,
Description: "",
AutoID: false,
Fields: []*schemapb.FieldSchema{
pk,
tz,
fVec,
},
Properties: []*commonpb.KeyValuePair{},
}
}
func constructCollectionSchemaWithAllType(
boolField, int32Field, int64Field, floatField, doubleField string,
floatVecField, binaryVecField, float16VecField, bfloat16VecField, structArrayField string,
geometryField string,
dim int,
collectionName string,
) *schemapb.CollectionSchema {
b := &schemapb.FieldSchema{
FieldID: 0,
Name: boolField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_Bool,
TypeParams: nil,
IndexParams: nil,
AutoID: false,
}
i32 := &schemapb.FieldSchema{
FieldID: 0,
Name: int32Field,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_Int32,
TypeParams: nil,
IndexParams: nil,
AutoID: false,
}
i64 := &schemapb.FieldSchema{
FieldID: 0,
Name: int64Field,
IsPrimaryKey: true,
Description: "",
DataType: schemapb.DataType_Int64,
TypeParams: nil,
IndexParams: nil,
AutoID: false,
}
f := &schemapb.FieldSchema{
FieldID: 0,
Name: floatField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_Float,
TypeParams: nil,
IndexParams: nil,
AutoID: false,
}
d := &schemapb.FieldSchema{
FieldID: 0,
Name: doubleField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_Double,
TypeParams: nil,
IndexParams: nil,
AutoID: false,
}
fVec := &schemapb.FieldSchema{
FieldID: 0,
Name: floatVecField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(dim),
},
},
IndexParams: nil,
AutoID: false,
}
bVec := &schemapb.FieldSchema{
FieldID: 0,
Name: binaryVecField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_BinaryVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(dim),
},
},
IndexParams: nil,
AutoID: false,
}
f16Vec := &schemapb.FieldSchema{
FieldID: 0,
Name: float16VecField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_Float16Vector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(dim),
},
},
IndexParams: nil,
AutoID: false,
}
bf16Vec := &schemapb.FieldSchema{
FieldID: 0,
Name: bfloat16VecField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_BFloat16Vector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(dim),
},
},
IndexParams: nil,
AutoID: false,
}
g := &schemapb.FieldSchema{
FieldID: 0,
Name: geometryField,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_Geometry,
TypeParams: nil,
IndexParams: nil,
AutoID: false,
}
// StructArrayField schema for testing
structArrayFields := []*schemapb.StructArrayFieldSchema{
{
FieldID: 0,
Name: structArrayField,
Description: "test struct array field",
Fields: []*schemapb.FieldSchema{
{
FieldID: 0,
Name: "sub_varchar_array",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxLengthKey,
Value: strconv.Itoa(testMaxVarCharLength),
},
{
Key: common.MaxCapacityKey,
Value: "100",
},
},
},
{
FieldID: 0,
Name: "sub_vector_array",
DataType: schemapb.DataType_ArrayOfVector,
ElementType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(dim),
},
{
Key: common.MaxCapacityKey,
Value: "10",
},
},
},
},
},
}
schema := &schemapb.CollectionSchema{
Name: collectionName,
Description: "",
AutoID: false,
StructArrayFields: structArrayFields,
}
if enableMultipleVectorFields {
schema.Fields = []*schemapb.FieldSchema{
b,
i32,
i64,
f,
d,
fVec,
bVec,
f16Vec,
bf16Vec,
g,
}
} else {
schema.Fields = []*schemapb.FieldSchema{
b,
i32,
i64,
f,
d,
fVec,
// bVec,
g,
}
}
return schema
}
func TestAlterCollection_AllowInsertAutoID_Validation(t *testing.T) {
ctx := context.Background()
buildRoot := func(autoID bool) *mocks.MockMixCoordClient {
root := mocks.NewMockMixCoordClient(t)
// InitMetaCache requires ListPolicy
root.EXPECT().ListPolicy(mock.Anything, mock.Anything, mock.Anything).Return(&internalpb.ListPolicyResponse{Status: merr.Success()}, nil).Once()
// Meta cache update path fetches partitions info
root.EXPECT().ShowPartitions(mock.Anything, mock.Anything, mock.Anything).Return(&milvuspb.ShowPartitionsResponse{Status: merr.Success()}, nil).Maybe()
// DescribeCollection returns a schema with PK autoID configurable
root.EXPECT().DescribeCollection(mock.Anything, mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
CollectionID: 1,
DbName: dbName,
Schema: &schemapb.CollectionSchema{
Name: "allow_autoid_test",
AutoID: false,
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", IsPrimaryKey: true, DataType: schemapb.DataType_Int64, AutoID: autoID},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "8"}}},
},
},
}, nil).Maybe()
return root
}
t.Run("success when PK autoID=true and allow_insert_autoid=true", func(t *testing.T) {
cache := globalMetaCache
defer func() { globalMetaCache = cache }()
root := buildRoot(true)
err := InitMetaCache(ctx, root)
assert.NoError(t, err)
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_AlterCollectionField},
DbName: dbName,
CollectionName: "allow_autoid_test",
Properties: []*commonpb.KeyValuePair{
{Key: common.AllowInsertAutoIDKey, Value: "true"},
},
},
}
err = task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("error when PK autoID=false and allow_insert_autoid=true", func(t *testing.T) {
cache := globalMetaCache
defer func() { globalMetaCache = cache }()
root := buildRoot(false)
err := InitMetaCache(ctx, root)
assert.NoError(t, err)
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_AlterCollectionField},
DbName: dbName,
CollectionName: "allow_autoid_test",
Properties: []*commonpb.KeyValuePair{
{Key: common.AllowInsertAutoIDKey, Value: "true"},
},
},
}
err = task.PreExecute(ctx)
assert.Error(t, err)
})
}
func constructPlaceholderGroup(
nq, dim int,
) *commonpb.PlaceholderGroup {
values := make([][]byte, 0, nq)
for i := 0; i < nq; i++ {
bs := make([]byte, 0, dim*4)
for j := 0; j < dim; j++ {
var buffer bytes.Buffer
f := rand.Float32()
err := binary.Write(&buffer, common.Endian, f)
if err != nil {
panic(err)
}
bs = append(bs, buffer.Bytes()...)
}
values = append(values, bs)
}
return &commonpb.PlaceholderGroup{
Placeholders: []*commonpb.PlaceholderValue{
{
Tag: "$0",
Type: commonpb.PlaceholderType_FloatVector,
Values: values,
},
},
}
}
func constructSearchRequest(
dbName, collectionName string,
expr string,
floatVecField string,
nq, dim, nprobe, topk, roundDecimal int,
) *milvuspb.SearchRequest {
params := make(map[string]string)
params["nprobe"] = strconv.Itoa(nprobe)
b, err := json.Marshal(params)
if err != nil {
panic(err)
}
plg := constructPlaceholderGroup(nq, dim)
plgBs, err := proto.Marshal(plg)
if err != nil {
panic(err)
}
return &milvuspb.SearchRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName,
PartitionNames: nil,
Dsl: expr,
SearchInput: &milvuspb.SearchRequest_PlaceholderGroup{
PlaceholderGroup: plgBs,
},
DslType: commonpb.DslType_BoolExprV1,
OutputFields: nil,
SearchParams: []*commonpb.KeyValuePair{
{
Key: common.MetricTypeKey,
Value: metric.L2,
},
{
Key: ParamsKey,
Value: string(b),
},
{
Key: AnnsFieldKey,
Value: floatVecField,
},
{
Key: TopKKey,
Value: strconv.Itoa(topk),
},
{
Key: RoundDecimalKey,
Value: strconv.Itoa(roundDecimal),
},
},
TravelTimestamp: 0,
GuaranteeTimestamp: 0,
}
}
func TestTranslateOutputFields(t *testing.T) {
const (
idFieldName = "id"
tsFieldName = "timestamp"
floatVectorFieldName = "float_vector"
binaryVectorFieldName = "binary_vector"
float16VectorFieldName = "float16_vector"
bfloat16VectorFieldName = "bfloat16_vector"
sparseFloatVectorFieldName = "sparse_float_vector"
)
var outputFields []string
var userOutputFields []string
var userDynamicFields []string
var requestedPK bool
var err error
collSchema := &schemapb.CollectionSchema{
Name: "TestTranslateOutputFields",
Description: "TestTranslateOutputFields",
AutoID: false,
Fields: []*schemapb.FieldSchema{
{Name: idFieldName, FieldID: 0, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{Name: tsFieldName, FieldID: 1, DataType: schemapb.DataType_Int64},
{Name: floatVectorFieldName, FieldID: 100, DataType: schemapb.DataType_FloatVector},
{Name: binaryVectorFieldName, FieldID: 101, DataType: schemapb.DataType_BinaryVector},
{Name: float16VectorFieldName, FieldID: 102, DataType: schemapb.DataType_Float16Vector},
{Name: bfloat16VectorFieldName, FieldID: 103, DataType: schemapb.DataType_BFloat16Vector},
{Name: sparseFloatVectorFieldName, FieldID: 104, DataType: schemapb.DataType_SparseFloatVector, IsFunctionOutput: true},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "bm25",
Type: schemapb.FunctionType_BM25,
OutputFieldNames: []string{sparseFloatVectorFieldName},
// omit other fields for brevity
},
},
}
schema := newSchemaInfo(collSchema)
// Test empty output fields
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{}, schema, false)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{}, outputFields)
assert.ElementsMatch(t, []string{}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.False(t, requestedPK)
// Test single field
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName}, schema, false)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{idFieldName}, outputFields)
assert.ElementsMatch(t, []string{idFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
// Test multiple fields
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName}, schema, false)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{idFieldName, tsFieldName}, outputFields)
assert.ElementsMatch(t, []string{idFieldName, tsFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
// Test with vector field
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName, floatVectorFieldName}, schema, false)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName}, outputFields)
assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
// Test without id field
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{tsFieldName, floatVectorFieldName}, schema, false)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, outputFields)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.False(t, requestedPK)
// Test wildcard - should not include function output fields
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*"}, schema, false)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
// Test wildcard with spaces
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{" * "}, schema, false)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
assert.ElementsMatch(t, []string{idFieldName, tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
_, _, _, _, _, err = translateOutputFields([]string{"*", sparseFloatVectorFieldName}, schema, false)
assert.Error(t, err)
_, _, _, _, _, err = translateOutputFields([]string{sparseFloatVectorFieldName}, schema, false)
assert.Error(t, err)
_, _, _, _, _, err = translateOutputFields([]string{sparseFloatVectorFieldName}, schema, true)
assert.Error(t, err)
// Test with removePkField=true
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{}, schema, true)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{}, outputFields)
assert.ElementsMatch(t, []string{}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.False(t, requestedPK)
// if removePkField is true, pk field should be removed from output fields
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*"}, schema, true)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName}, schema, true)
assert.Equal(t, nil, err)
assert.ElementsMatch(t, []string{tsFieldName}, outputFields)
assert.ElementsMatch(t, []string{tsFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{idFieldName, tsFieldName, floatVectorFieldName}, schema, true)
assert.Equal(t, nil, err)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, outputFields)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*"}, schema, true)
assert.Equal(t, nil, err)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*", tsFieldName}, schema, true)
assert.Equal(t, nil, err)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"*", floatVectorFieldName}, schema, true)
assert.Equal(t, nil, err)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, outputFields)
assert.ElementsMatch(t, []string{tsFieldName, floatVectorFieldName, binaryVectorFieldName, float16VectorFieldName, bfloat16VectorFieldName}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.True(t, requestedPK)
// Test non-existent field, dynamic field not enabled
_, _, _, _, _, err = translateOutputFields([]string{"A"}, schema, true)
assert.Error(t, err)
t.Run("enable dynamic schema", func(t *testing.T) {
collSchema := &schemapb.CollectionSchema{
Name: "TestTranslateOutputFields",
Description: "TestTranslateOutputFields",
AutoID: false,
EnableDynamicField: true,
Fields: []*schemapb.FieldSchema{
{Name: idFieldName, FieldID: 1, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{Name: tsFieldName, FieldID: 2, DataType: schemapb.DataType_Int64},
{Name: floatVectorFieldName, FieldID: 100, DataType: schemapb.DataType_FloatVector},
{Name: binaryVectorFieldName, FieldID: 101, DataType: schemapb.DataType_BinaryVector},
{Name: common.MetaFieldName, FieldID: 102, DataType: schemapb.DataType_JSON, IsDynamic: true},
},
}
schema := newSchemaInfo(collSchema)
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"A", idFieldName}, schema, true)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{common.MetaFieldName}, outputFields)
assert.ElementsMatch(t, []string{"A"}, userOutputFields)
assert.ElementsMatch(t, []string{"A"}, userDynamicFields)
assert.True(t, requestedPK)
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"$meta[\"A\"]", idFieldName}, schema, true)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{common.MetaFieldName}, outputFields)
assert.ElementsMatch(t, []string{"$meta[\"A\"]"}, userOutputFields)
assert.ElementsMatch(t, []string{"A"}, userDynamicFields)
assert.True(t, requestedPK)
// Test invalid dynamic field expressions
_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, `$meta["A"]["B"]`}, schema, true)
assert.Error(t, err)
_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "$meta[\"\"]"}, schema, true)
assert.Error(t, err)
_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "$meta[]"}, schema, true)
assert.Error(t, err)
_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "$meta["}, schema, true)
assert.Error(t, err)
_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "[]"}, schema, true)
assert.Error(t, err)
_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, "A > 1"}, schema, true)
assert.Error(t, err)
_, _, _, _, _, err = translateOutputFields([]string{idFieldName, floatVectorFieldName, ""}, schema, true)
assert.Error(t, err)
})
}
func TestTranslateOutputFields_StructArrayField(t *testing.T) {
const (
idFieldName = "id"
tsFieldName = "timestamp"
floatVectorFieldName = "float_vector"
binaryVectorFieldName = "binary_vector"
float16VectorFieldName = "float16_vector"
bfloat16VectorFieldName = "bfloat16_vector"
sparseFloatVectorFieldName = "sparse_float_vector"
)
var outputFields []string
var userOutputFields []string
var userDynamicFields []string
var requestedPK bool
var err error
collSchema := &schemapb.CollectionSchema{
Name: "TestTranslateOutputFields",
Description: "TestTranslateOutputFields",
AutoID: false,
Fields: []*schemapb.FieldSchema{
{Name: idFieldName, FieldID: 0, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{Name: tsFieldName, FieldID: 1, DataType: schemapb.DataType_Int64},
{Name: floatVectorFieldName, FieldID: 100, DataType: schemapb.DataType_FloatVector},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 101,
Name: "struct_array_field",
Fields: []*schemapb.FieldSchema{
{Name: "sub_field", FieldID: 102, DataType: schemapb.DataType_Array, ElementType: schemapb.DataType_Int64},
{Name: "sub_vector_field", FieldID: 103, DataType: schemapb.DataType_ArrayOfVector, ElementType: schemapb.DataType_FloatVector},
},
},
},
}
schema := newSchemaInfo(collSchema)
// Test struct array field
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"sub_vector_field"}, schema, false)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{"sub_vector_field"}, outputFields)
assert.ElementsMatch(t, []string{"sub_vector_field"}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.False(t, requestedPK)
outputFields, userOutputFields, userDynamicFields, _, requestedPK, err = translateOutputFields([]string{"struct_array_field", "sub_field"}, schema, false)
assert.NoError(t, err)
assert.ElementsMatch(t, []string{"sub_vector_field", "sub_field"}, outputFields)
assert.ElementsMatch(t, []string{"sub_vector_field", "sub_field"}, userOutputFields)
assert.ElementsMatch(t, []string{}, userDynamicFields)
assert.False(t, requestedPK)
}
func TestAddFieldTask(t *testing.T) {
rc := NewMixCoordMock()
ctx := context.Background()
prefix := "TestAddFieldTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
collectionID := int64(1)
int64Field := "int64"
floatVecField := "fvec"
varCharField := "varChar"
rc.collName2ID[collectionName] = collectionID
rc.collID2Meta[collectionID] = collectionMeta{
name: collectionName,
id: collectionID,
schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int64, AutoID: true, Name: "ID"},
{
FieldID: 101, DataType: schemapb.DataType_FloatVector, Name: "vector",
TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
},
},
},
},
}
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type[int64Field] = schemapb.DataType_Int64
fieldName2Type[varCharField] = schemapb.DataType_VarChar
fieldName2Type[floatVecField] = schemapb.DataType_FloatVector
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, int64Field, false)
fSchema := &schemapb.FieldSchema{
Name: "add",
DataType: schemapb.DataType_Bool,
Nullable: true,
}
bytes, err := proto.Marshal(fSchema)
assert.NoError(t, err)
task := &addCollectionFieldTask{
Condition: NewTaskCondition(ctx),
AddCollectionFieldRequest: &milvuspb.AddCollectionFieldRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName,
Schema: bytes,
},
ctx: ctx,
mixCoord: rc,
result: nil,
oldSchema: schema,
}
t.Run("on enqueue", func(t *testing.T) {
err := task.OnEnqueue()
assert.NoError(t, err)
assert.Equal(t, commonpb.MsgType_AddCollectionField, task.Type())
})
t.Run("ctx", func(t *testing.T) {
traceCtx := task.TraceCtx()
assert.NotNil(t, traceCtx)
})
t.Run("id", func(t *testing.T) {
id := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
task.SetID(id)
assert.Equal(t, id, task.ID())
})
t.Run("name", func(t *testing.T) {
assert.Equal(t, AddFieldTaskName, task.Name())
})
t.Run("ts", func(t *testing.T) {
ts := Timestamp(time.Now().UnixNano())
task.SetTs(ts)
assert.Equal(t, ts, task.BeginTs())
assert.Equal(t, ts, task.EndTs())
})
t.Run("process task", func(t *testing.T) {
var err error
// nil collection schema
task.oldSchema = nil
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
bytes, err := proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
task.oldSchema = schema
err = task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.ErrorCode)
err = task.PostExecute(ctx)
assert.NoError(t, err)
})
t.Run("PreExecute", func(t *testing.T) {
var err error
err = task.PreExecute(ctx)
assert.NoError(t, err)
// nil schema
task.Schema = nil
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
// too many fields
Params.Save(Params.ProxyCfg.MaxFieldNum.Key, fmt.Sprint(task.oldSchema.Fields))
fSchema := &schemapb.FieldSchema{
Name: "add_field",
}
bytes, err := proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
Params.Reset(Params.ProxyCfg.MaxFieldNum.Key)
// invalid field type
fSchema = &schemapb.FieldSchema{
DataType: schemapb.DataType_None,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
fSchema = &schemapb.FieldSchema{
Name: "vec_field",
DataType: schemapb.DataType_FloatVector,
Nullable: true,
TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
},
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.NoError(t, err)
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
})
fSchema = &schemapb.FieldSchema{
Name: "text_field",
DataType: schemapb.DataType_Text,
Nullable: true,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.NoError(t, err)
fSchema = &schemapb.FieldSchema{
Name: "text_not_nullable",
DataType: schemapb.DataType_Text,
Nullable: false,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "added field must be nullable")
fSchema = &schemapb.FieldSchema{
Name: "text_default",
DataType: schemapb.DataType_Text,
Nullable: true,
DefaultValue: &schemapb.ValueField{
Data: &schemapb.ValueField_StringData{StringData: "default text"},
},
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "default value is not supported when adding TEXT field")
fSchema = &schemapb.FieldSchema{
Name: "sparse_vec",
DataType: schemapb.DataType_SparseFloatVector,
Nullable: true,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.NoError(t, err)
fSchema = &schemapb.FieldSchema{
Name: "array_struct",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Struct,
Nullable: true,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "4"},
},
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "element type Struct is not supported")
// not support system field
fSchema = &schemapb.FieldSchema{
Name: common.TimeStampFieldName,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
// must be nullable
fSchema = &schemapb.FieldSchema{
Nullable: false,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
// not support pk field
fSchema = &schemapb.FieldSchema{
IsPrimaryKey: true,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
// not support add partition key
fSchema = &schemapb.FieldSchema{
IsPartitionKey: true,
DataType: schemapb.DataType_Bool,
Nullable: true,
Name: "new field",
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
Params.Save(Params.ProxyCfg.MustUsePartitionKey.Key, "true")
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
Params.Reset(Params.ProxyCfg.MustUsePartitionKey.Key)
// not support autoID
fSchema = &schemapb.FieldSchema{
AutoID: true,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
// more ClusteringKey field — use a valid type so the duplicate-key check is the gate
fSchema = &schemapb.FieldSchema{
Name: "ck_field",
DataType: schemapb.DataType_Int64,
Nullable: true,
IsClusteringKey: true,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
task.oldSchema = schema
task.oldSchema.Fields = append(task.oldSchema.Fields, &schemapb.FieldSchema{
IsClusteringKey: true,
})
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
// unsupported clustering key type (JSON, Bool, Array)
// Use a fresh schema (no existing clustering key) so the type check is the only gate
freshSchema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, int64Field, false)
task.oldSchema = freshSchema
for _, unsupportedType := range []schemapb.DataType{
schemapb.DataType_JSON,
schemapb.DataType_Bool,
schemapb.DataType_Array,
} {
fSchema = &schemapb.FieldSchema{
Name: "ck_field",
DataType: unsupportedType,
Nullable: true,
IsClusteringKey: true,
}
if unsupportedType == schemapb.DataType_Array {
fSchema.ElementType = schemapb.DataType_Int64
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err, "expected error for unsupported clustering key type %v", unsupportedType)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
}
// valid clustering key type (Int64) should be accepted
fSchema = &schemapb.FieldSchema{
Name: "valid_ck",
DataType: schemapb.DataType_Int64,
Nullable: true,
IsClusteringKey: true,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
task.oldSchema = freshSchema // no existing clustering key
err = task.PreExecute(ctx)
assert.NoError(t, err)
// fieldName invalid
fSchema = &schemapb.FieldSchema{
Name: "",
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
// duplicated FieldName
fSchema = &schemapb.FieldSchema{
Name: varCharField,
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
oldSchemaWithVectors := constructCollectionSchemaByDataType(collectionName, fieldName2Type, int64Field, false)
for i := 0; i < Params.ProxyCfg.MaxVectorFieldNum.GetAsInt(); i++ {
oldSchemaWithVectors.Fields = append(oldSchemaWithVectors.Fields, &schemapb.FieldSchema{
Name: fmt.Sprintf("vec_%d", i),
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
},
})
}
task.oldSchema = oldSchemaWithVectors
fSchema = &schemapb.FieldSchema{
Name: "new_vec_field",
DataType: schemapb.DataType_FloatVector,
Nullable: true,
TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
},
}
bytes, err = proto.Marshal(fSchema)
assert.NoError(t, err)
task.Schema = bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "maximum vector field's number should be limited to")
})
}
func TestAddCollectionStructFieldTaskPreExecute(t *testing.T) {
ctx := context.Background()
oldSchema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "pk",
DataType: schemapb.DataType_Int64,
IsPrimaryKey: true,
},
},
}
structField := newAddStructFieldSchema("profile")
task := &addCollectionStructFieldTask{
Condition: NewTaskCondition(ctx),
AddCollectionStructFieldRequest: &milvuspb.AddCollectionStructFieldRequest{
DbName: "",
CollectionName: oldSchema.GetName(),
StructArrayFieldSchema: structField,
},
ctx: ctx,
mixCoord: NewMixCoordMock(),
oldSchema: oldSchema,
}
require.NoError(t, task.OnEnqueue())
assert.Equal(t, commonpb.MsgType_AddCollectionField, task.Type())
require.NoError(t, task.PreExecute(ctx))
got := task.GetStructArrayFieldSchema()
require.NotSame(t, structField, got)
require.Len(t, got.GetFields(), 2)
assert.Equal(t, "profile[ints]", got.GetFields()[0].GetName())
assert.Equal(t, "profile[vectors]", got.GetFields()[1].GetName())
assert.True(t, got.GetFields()[0].GetNullable())
assert.True(t, got.GetFields()[1].GetNullable())
// PreExecute works on a cloned schema, so callers do not observe partial mutation on failure.
assert.Equal(t, "ints", structField.GetFields()[0].GetName())
assert.Equal(t, "vectors", structField.GetFields()[1].GetName())
}
func TestValidateAddStructFieldRequest(t *testing.T) {
baseSchema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{
Name: "pk",
DataType: schemapb.DataType_Int64,
},
},
}
t.Run("valid", func(t *testing.T) {
err := validateAddStructFieldRequest(proto.Clone(baseSchema).(*schemapb.CollectionSchema), newAddStructFieldSchema("profile"))
require.NoError(t, err)
})
t.Run("struct must be nullable", func(t *testing.T) {
structField := newAddStructFieldSchema("profile")
structField.Nullable = false
err := validateAddStructFieldRequest(proto.Clone(baseSchema).(*schemapb.CollectionSchema), structField)
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "added struct field must be nullable")
})
t.Run("duplicate parent name", func(t *testing.T) {
err := validateAddStructFieldRequest(proto.Clone(baseSchema).(*schemapb.CollectionSchema), newAddStructFieldSchema("pk"))
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "duplicated field name pk")
})
t.Run("duplicate transformed sub field name", func(t *testing.T) {
schema := proto.Clone(baseSchema).(*schemapb.CollectionSchema)
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
Name: "profile[ints]",
DataType: schemapb.DataType_Int64,
})
err := validateAddStructFieldRequest(schema, newAddStructFieldSchema("profile"))
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "duplicated field name profile[ints]")
})
t.Run("max field count includes struct parent", func(t *testing.T) {
Params.Save(Params.ProxyCfg.MaxFieldNum.Key, "3")
defer Params.Reset(Params.ProxyCfg.MaxFieldNum.Key)
err := validateAddStructFieldRequest(proto.Clone(baseSchema).(*schemapb.CollectionSchema), newAddStructFieldSchema("profile"))
require.Error(t, err)
assert.Contains(t, err.Error(), "maximum field's number should be limited to 3")
})
t.Run("vector count includes array of vector sub field", func(t *testing.T) {
Params.Save(Params.ProxyCfg.MaxVectorFieldNum.Key, "0")
defer Params.Reset(Params.ProxyCfg.MaxVectorFieldNum.Key)
err := validateAddStructFieldRequest(proto.Clone(baseSchema).(*schemapb.CollectionSchema), newAddStructFieldSchema("profile"))
require.Error(t, err)
assert.Contains(t, err.Error(), "maximum vector field's number should be limited to 0")
})
}
func newAddStructFieldSchema(name string) *schemapb.StructArrayFieldSchema {
return &schemapb.StructArrayFieldSchema{
Name: name,
Nullable: true,
Fields: []*schemapb.FieldSchema{
{
Name: "ints",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "16"},
},
},
{
Name: "vectors",
DataType: schemapb.DataType_ArrayOfVector,
ElementType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "8"},
{Key: common.MaxCapacityKey, Value: "16"},
},
},
},
}
}
func TestCreateCollectionTask(t *testing.T) {
mix := NewMixCoordMock()
ctx := context.Background()
shardsNum := common.DefaultShardsNum
prefix := "TestCreateCollectionTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
int64Field := "int64"
floatVecField := "fvec"
varCharField := "varChar"
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type[int64Field] = schemapb.DataType_Int64
fieldName2Type[varCharField] = schemapb.DataType_VarChar
fieldName2Type[floatVecField] = schemapb.DataType_FloatVector
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, int64Field, false)
var marshaledSchema []byte
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: mix,
result: nil,
schema: nil,
}
t.Run("on enqueue", func(t *testing.T) {
err := task.OnEnqueue()
assert.NoError(t, err)
assert.Equal(t, commonpb.MsgType_CreateCollection, task.Type())
})
t.Run("ctx", func(t *testing.T) {
traceCtx := task.TraceCtx()
assert.NotNil(t, traceCtx)
})
t.Run("id", func(t *testing.T) {
id := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
task.SetID(id)
assert.Equal(t, id, task.ID())
})
t.Run("name", func(t *testing.T) {
assert.Equal(t, CreateCollectionTaskName, task.Name())
})
t.Run("ts", func(t *testing.T) {
ts := Timestamp(time.Now().UnixNano())
task.SetTs(ts)
assert.Equal(t, ts, task.BeginTs())
assert.Equal(t, ts, task.EndTs())
})
t.Run("process task", func(t *testing.T) {
var err error
err = task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.ErrorCode)
// recreate -> fail
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.ErrorCode)
err = task.PostExecute(ctx)
assert.NoError(t, err)
})
t.Run("PreExecute", func(t *testing.T) {
var err error
err = task.PreExecute(ctx)
assert.NoError(t, err)
Params.Save(Params.ProxyCfg.MustUsePartitionKey.Key, "true")
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterMissing)
Params.Reset(Params.ProxyCfg.MustUsePartitionKey.Key)
task.Schema = []byte{0x1, 0x2, 0x3, 0x4}
err = task.PreExecute(ctx)
assert.Error(t, err)
task.Schema = marshaledSchema
task.ShardsNum = Params.ProxyCfg.MaxShardNum.GetAsInt32() + 1
err = task.PreExecute(ctx)
assert.Error(t, err)
task.ShardsNum = shardsNum
reqBackup := proto.Clone(task.CreateCollectionRequest).(*milvuspb.CreateCollectionRequest)
schemaBackup := proto.Clone(schema).(*schemapb.CollectionSchema)
schemaWithTooManyFields := &schemapb.CollectionSchema{
Name: collectionName,
Description: "",
AutoID: false,
Fields: make([]*schemapb.FieldSchema, Params.ProxyCfg.MaxFieldNum.GetAsInt32()+1),
}
marshaledSchemaWithTooManyFields, err := proto.Marshal(schemaWithTooManyFields)
assert.NoError(t, err)
task.Schema = marshaledSchemaWithTooManyFields
err = task.PreExecute(ctx)
assert.Error(t, err)
// too many vector fields
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
schema.Fields = append(schema.Fields, schema.Fields[0])
for i := 0; i < Params.ProxyCfg.MaxVectorFieldNum.GetAsInt(); i++ {
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
FieldID: 101,
Name: floatVecField + "_" + strconv.Itoa(i),
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(testVecDim),
},
},
IndexParams: nil,
AutoID: false,
})
}
tooManyVectorFieldsSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = tooManyVectorFieldsSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
// without vector field
schema = &schemapb.CollectionSchema{
Name: collectionName,
Description: "",
AutoID: false,
Fields: []*schemapb.FieldSchema{
{
Name: "id",
DataType: schemapb.DataType_Int64,
IsPrimaryKey: true,
},
},
}
noVectorSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = noVectorSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
task.CreateCollectionRequest = reqBackup
// validateCollectionName
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
schema.Name = " " // empty
emptyNameSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = emptyNameSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
schema.Name = prefix
for i := 0; i < Params.ProxyCfg.MaxNameLength.GetAsInt(); i++ {
schema.Name += strconv.Itoa(i % 10)
}
tooLongNameSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = tooLongNameSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
schema.Description = strings.Repeat("a", Params.ProxyCfg.MaxCollectionDescriptionLength.GetAsInt()+1)
tooLongDescriptionSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = tooLongDescriptionSchema
err = task.PreExecute(ctx)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
schema.Name = "$" // invalid first char
invalidFirstCharSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = invalidFirstCharSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
// validateDuplicatedFieldName
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
schema.Fields = append(schema.Fields, schema.Fields[0])
duplicatedFieldsSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = duplicatedFieldsSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
// validatePrimaryKey
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
for idx := range schema.Fields {
schema.Fields[idx].IsPrimaryKey = false
}
noPrimaryFieldsSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = noPrimaryFieldsSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
// validateFieldName
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
for idx := range schema.Fields {
schema.Fields[idx].Name = "$"
}
invalidFieldNameSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = invalidFieldNameSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
// validateMaxLengthPerRow
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
for idx := range schema.Fields {
if schema.Fields[idx].DataType == schemapb.DataType_VarChar {
schema.Fields[idx].TypeParams = nil
}
}
noTypeParamsSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = noTypeParamsSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
// ValidateVectorField
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
for idx := range schema.Fields {
if schema.Fields[idx].DataType == schemapb.DataType_FloatVector ||
schema.Fields[idx].DataType == schemapb.DataType_BinaryVector {
schema.Fields[idx].TypeParams = nil
}
}
noDimSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = noDimSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
// ValidateVectorField
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
for _, field := range schema.Fields {
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.FieldSkipLoadKey,
Value: "true",
})
}
// Validate default load list
skipLoadSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = skipLoadSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
for idx := range schema.Fields {
if schema.Fields[idx].DataType == schemapb.DataType_FloatVector ||
schema.Fields[idx].DataType == schemapb.DataType_BinaryVector {
schema.Fields[idx].TypeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "not int",
},
}
}
}
dimNotIntSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = dimNotIntSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
for idx := range schema.Fields {
if schema.Fields[idx].DataType == schemapb.DataType_FloatVector ||
schema.Fields[idx].DataType == schemapb.DataType_BinaryVector {
schema.Fields[idx].TypeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(Params.ProxyCfg.MaxDimension.GetAsInt() + 1),
},
}
}
}
tooLargeDimSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = tooLargeDimSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
schema.Fields[1].DataType = schemapb.DataType_BinaryVector
schema.Fields[1].TypeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(Params.ProxyCfg.MaxDimension.GetAsInt() + 1),
},
}
binaryTooLargeDimSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = binaryTooLargeDimSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
schema = proto.Clone(schemaBackup).(*schemapb.CollectionSchema)
lastFieldID := schema.Fields[len(schema.Fields)-1].FieldID
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
FieldID: lastFieldID + 1,
Name: "second_vector",
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(128),
},
},
IndexParams: nil,
AutoID: false,
})
twoVecFieldsSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = twoVecFieldsSchema
err = task.PreExecute(ctx)
if enableMultipleVectorFields {
assert.NoError(t, err)
} else {
assert.Error(t, err)
}
task.CreateCollectionRequest = reqBackup
// Test StructArrayField validation
structArrayFieldSchema := constructCollectionSchemaWithStructArrayField(collectionName+"_struct", testStructArrayField, false)
// Test valid StructArrayField
validStructSchema, err := proto.Marshal(structArrayFieldSchema)
assert.NoError(t, err)
task.Schema = validStructSchema
err = task.PreExecute(ctx)
assert.NoError(t, err)
// Test invalid StructArrayField - empty fields
invalidStructSchema := proto.Clone(structArrayFieldSchema).(*schemapb.CollectionSchema)
invalidStructSchema.StructArrayFields[0].Fields = []*schemapb.FieldSchema{} // Empty fields
invalidStructSchemaBytes, err := proto.Marshal(invalidStructSchema)
assert.NoError(t, err)
task.Schema = invalidStructSchemaBytes
err = task.PreExecute(ctx)
assert.Error(t, err)
// Test invalid StructArrayField - invalid field name
invalidStructSchema2 := proto.Clone(structArrayFieldSchema).(*schemapb.CollectionSchema)
invalidStructSchema2.StructArrayFields[0].Fields[0].Name = "$invalid" // Invalid field name
invalidStructSchema2Bytes, err := proto.Marshal(invalidStructSchema2)
assert.NoError(t, err)
task.Schema = invalidStructSchema2Bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
// Test invalid StructArrayField - non-array/non-vector-array sub-field
invalidStructSchema3 := proto.Clone(structArrayFieldSchema).(*schemapb.CollectionSchema)
invalidStructSchema3.StructArrayFields[0].Fields[0].DataType = schemapb.DataType_Int64 // Should be Array or ArrayOfVector
invalidStructSchema3Bytes, err := proto.Marshal(invalidStructSchema3)
assert.NoError(t, err)
task.Schema = invalidStructSchema3Bytes
err = task.PreExecute(ctx)
assert.Error(t, err)
ttlFieldSchema := constructCollectionSchemaWithTTLField(collectionName+"_ttl", testTTLField)
// Test invalid ttl field - ttl field not found
ttlFieldSchemaBytes, err := proto.Marshal(ttlFieldSchema)
assert.NoError(t, err)
task.Properties = make([]*commonpb.KeyValuePair, 0)
task.Schema = ttlFieldSchemaBytes
task.Properties = append(task.Properties, &commonpb.KeyValuePair{
Key: common.CollectionTTLFieldKey,
Value: "invalid",
})
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
// Test invalid ttl field - exist with collection ttl
task.Properties = make([]*commonpb.KeyValuePair, 0)
task.Properties = append(task.Properties, &commonpb.KeyValuePair{
Key: common.CollectionTTLFieldKey,
Value: testTTLField,
})
task.Properties = append(task.Properties, &commonpb.KeyValuePair{
Key: common.CollectionTTLConfigKey,
Value: "100",
})
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
// Test invalid ttl field - ttl field is not timestamptz
task.Properties = make([]*commonpb.KeyValuePair, 0)
task.Properties = append(task.Properties, &commonpb.KeyValuePair{
Key: common.CollectionTTLFieldKey,
Value: "pk",
})
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
// Restore original schema for remaining tests
task.CreateCollectionRequest = reqBackup
})
t.Run("specify dynamic field", func(t *testing.T) {
dynamicField := &schemapb.FieldSchema{
Name: "json",
IsDynamic: true,
}
var marshaledSchema []byte
schema2 := &schemapb.CollectionSchema{
Name: collectionName,
Fields: append(schema.Fields, dynamicField),
}
marshaledSchema, err := proto.Marshal(schema2)
assert.NoError(t, err)
task2 := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: mix,
result: nil,
schema: nil,
}
err = task2.OnEnqueue()
assert.NoError(t, err)
err = task2.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("collection with embedding function ", func(t *testing.T) {
paramtable.Init()
paramtable.Get().CredentialCfg.Credential.GetFunc = func() map[string]string {
return map[string]string{
"mock.apikey": "mock",
}
}
ts := embedding.CreateOpenAIEmbeddingServer()
defer ts.Close()
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
return map[string]string{
"openai.url": ts.URL,
}
}
schema.Functions = []*schemapb.FunctionSchema{
{
Name: "test",
Type: schemapb.FunctionType_TextEmbedding,
InputFieldNames: []string{varCharField},
OutputFieldNames: []string{floatVecField},
Params: []*commonpb.KeyValuePair{
{Key: "provider", Value: "openai"},
{Key: "model_name", Value: "text-embedding-ada-002"},
{Key: "credential", Value: "mock"},
{Key: "dim", Value: "128"},
},
},
}
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task2 := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: mix,
result: nil,
schema: nil,
}
err = task2.OnEnqueue()
assert.NoError(t, err)
err = task2.PreExecute(ctx)
assert.NoError(t, err)
})
}
func TestCreateCollectionTaskExternalCollection(t *testing.T) {
mix := NewMixCoordMock()
ctx := context.Background()
collectionName := "external_collection_" + funcutil.GenRandomStr()
buildExternalSchema := func() *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Name: collectionName,
ExternalSource: "s3://bucket/prefix",
ExternalSpec: `{"format":"parquet"}`,
Fields: []*schemapb.FieldSchema{
{
FieldID: 1,
Name: "text_field",
DataType: schemapb.DataType_VarChar,
ExternalField: "text_col",
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "64"},
},
},
{
FieldID: 2,
Name: "vec_field",
DataType: schemapb.DataType_FloatVector,
ExternalField: "vec_col",
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "16"},
},
},
},
}
}
marshal := func(schema *schemapb.CollectionSchema) []byte {
bytes, err := proto.Marshal(schema)
require.NoError(t, err)
return bytes
}
task := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
CollectionName: collectionName,
Schema: marshal(buildExternalSchema()),
ShardsNum: common.DefaultShardsNum,
},
ctx: ctx,
mixCoord: mix,
}
t.Run("valid external schema", func(t *testing.T) {
err := task.OnEnqueue()
require.NoError(t, err)
err = task.PreExecute(ctx)
require.NoError(t, err)
})
t.Run("source spec without external field rejected", func(t *testing.T) {
schema := buildExternalSchema()
for _, field := range schema.GetFields() {
field.ExternalField = ""
}
freshTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
CollectionName: collectionName,
Schema: marshal(schema),
ShardsNum: common.DefaultShardsNum,
},
ctx: ctx,
mixCoord: mix,
}
err := freshTask.OnEnqueue()
require.NoError(t, err)
err = freshTask.PreExecute(ctx)
require.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "external_field mappings")
})
t.Run("virtual PK injected after PreExecute", func(t *testing.T) {
schema := buildExternalSchema()
freshTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
CollectionName: collectionName,
Schema: marshal(schema),
ShardsNum: common.DefaultShardsNum,
},
ctx: ctx,
mixCoord: mix,
}
err := freshTask.OnEnqueue()
require.NoError(t, err)
err = freshTask.PreExecute(ctx)
require.NoError(t, err)
// After PreExecute, schema should have virtual PK prepended
assert.Equal(t, common.VirtualPKFieldName, freshTask.schema.Fields[0].Name)
assert.True(t, freshTask.schema.Fields[0].IsPrimaryKey)
assert.True(t, freshTask.schema.Fields[0].AutoID)
assert.Equal(t, schemapb.DataType_Int64, freshTask.schema.Fields[0].DataType)
// Original fields should follow
assert.Equal(t, "text_field", freshTask.schema.Fields[1].Name)
assert.Equal(t, "vec_field", freshTask.schema.Fields[2].Name)
})
t.Run("dynamic field forbidden", func(t *testing.T) {
schema := buildExternalSchema()
schema.EnableDynamicField = true
task.Schema = marshal(schema)
err := task.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("partition key forbidden", func(t *testing.T) {
schema := buildExternalSchema()
schema.Fields[0].IsPartitionKey = true
task.Schema = marshal(schema)
err := task.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("clustering key forbidden", func(t *testing.T) {
schema := buildExternalSchema()
schema.Fields[0].IsClusteringKey = true
task.Schema = marshal(schema)
err := task.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("auto id forbidden", func(t *testing.T) {
schema := buildExternalSchema()
schema.Fields[0].AutoID = true
task.Schema = marshal(schema)
err := task.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("text match allowed", func(t *testing.T) {
schema := buildExternalSchema()
schema.Fields[0].TypeParams = append(schema.Fields[0].TypeParams,
&commonpb.KeyValuePair{
Key: "enable_match",
Value: "true",
},
&commonpb.KeyValuePair{
Key: common.EnableAnalyzerKey,
Value: "true",
})
task.Schema = marshal(schema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("struct field forbidden", func(t *testing.T) {
schema := buildExternalSchema()
schema.StructArrayFields = []*schemapb.StructArrayFieldSchema{
{
FieldID: 3,
Name: "struct_array_field",
Fields: []*schemapb.FieldSchema{
{
FieldID: 4,
Name: "nested_array",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "10"},
},
},
},
},
}
task.Schema = marshal(schema)
err := task.PreExecute(ctx)
assert.Error(t, err)
})
}
func TestHasCollectionTask(t *testing.T) {
mixc := NewMixCoordMock()
defer mixc.Close()
ctx := context.Background()
InitMetaCache(ctx, mixc)
prefix := "TestHasCollectionTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
shardsNum := int32(2)
int64Field := "int64"
floatVecField := "fvec"
dim := 128
schema := constructCollectionSchema(int64Field, floatVecField, dim, collectionName)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DropCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
}
// CreateCollection
task := &hasCollectionTask{
Condition: NewTaskCondition(ctx),
HasCollectionRequest: &milvuspb.HasCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_HasCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
},
ctx: ctx,
mixCoord: mixc,
result: nil,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_HasCollection, task.Type())
assert.Equal(t, UniqueID(100), task.ID())
assert.Equal(t, Timestamp(100), task.BeginTs())
assert.Equal(t, Timestamp(100), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.GetBase().GetSourceID())
// missing collectionID in globalMetaCache
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, false, task.result.Value)
// createIsoCollection in RootCood and fill GlobalMetaCache
mixc.CreateCollection(ctx, createColReq)
globalMetaCache.GetCollectionID(ctx, GetCurDBNameFromContextOrDefault(ctx), collectionName)
// success to drop collection
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, true, task.result.Value)
// illegal name
task.CollectionName = "#0xc0de"
err = task.PreExecute(ctx)
assert.Error(t, err)
task.CollectionName = collectionName
// invalidate collection cache, trigger rootcoord rpc
globalMetaCache.RemoveCollection(ctx, dbName, collectionName, 0)
// rc return collection not found error
mixc.describeCollectionFunc = func(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
return nil, merr.WrapErrCollectionNotFoundWithDB(dbName, collectionName)
}
err = task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
assert.False(t, task.result.GetValue())
assert.NotNil(t, task.result.GetStatus())
// rootcoord failed to get response
mixc.updateState(commonpb.StateCode_Abnormal)
err = task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.Error(t, err)
}
func TestDescribeCollectionTask(t *testing.T) {
mixc := NewMixCoordMock()
defer mixc.Close()
ctx := context.Background()
InitMetaCache(ctx, mixc)
prefix := "TestDescribeCollectionTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
// CreateCollection
task := &describeCollectionTask{
Condition: NewTaskCondition(ctx),
DescribeCollectionRequest: &milvuspb.DescribeCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DescribeCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
},
ctx: ctx,
mixCoord: mixc,
result: nil,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_DescribeCollection, task.Type())
assert.Equal(t, UniqueID(100), task.ID())
assert.Equal(t, Timestamp(100), task.BeginTs())
assert.Equal(t, Timestamp(100), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.GetBase().GetSourceID())
// missing collectionID in globalMetaCache
err := task.Execute(ctx)
assert.NoError(t, err)
// illegal name
task.CollectionName = "#0xc0de"
err = task.PreExecute(ctx)
assert.Error(t, err)
// describe collection with id
task.CollectionID = 1
task.CollectionName = ""
err = task.PreExecute(ctx)
assert.NoError(t, err)
task.CollectionID = 0
task.CollectionName = collectionName
err = task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, task.result.GetStatus().GetErrorCode())
}
func TestDescribeCollectionTask_ShardsNum1(t *testing.T) {
mix := NewMixCoordMock()
ctx := context.Background()
InitMetaCache(ctx, mix)
prefix := "TestDescribeCollectionTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
shardsNum := common.DefaultShardsNum
int64Field := "int64"
floatVecField := "fvec"
dim := 128
schema := constructCollectionSchema(int64Field, floatVecField, dim, collectionName)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DropCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
}
mix.CreateCollection(ctx, createColReq)
globalMetaCache.GetCollectionID(ctx, GetCurDBNameFromContextOrDefault(ctx), collectionName)
// CreateCollection
task := &describeCollectionTask{
Condition: NewTaskCondition(ctx),
DescribeCollectionRequest: &milvuspb.DescribeCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DescribeCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
},
ctx: ctx,
mixCoord: mix,
result: nil,
}
err = task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.GetStatus().GetErrorCode())
assert.Equal(t, shardsNum, task.result.ShardsNum)
assert.Equal(t, collectionName, task.result.GetCollectionName())
}
func TestDescribeCollectionTask_EnableDynamicSchema(t *testing.T) {
mix := NewMixCoordMock()
ctx := context.Background()
InitMetaCache(ctx, mix)
prefix := "TestDescribeCollectionTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
shardsNum := common.DefaultShardsNum
int64Field := "int64"
floatVecField := "fvec"
dim := 128
schema := constructCollectionSchemaEnableDynamicSchema(int64Field, floatVecField, dim, collectionName)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DropCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
}
mix.CreateCollection(ctx, createColReq)
globalMetaCache.GetCollectionID(ctx, dbName, collectionName)
// CreateCollection
task := &describeCollectionTask{
Condition: NewTaskCondition(ctx),
DescribeCollectionRequest: &milvuspb.DescribeCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DescribeCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
},
ctx: ctx,
mixCoord: mix,
result: nil,
}
err = task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.GetStatus().GetErrorCode())
assert.Equal(t, shardsNum, task.result.ShardsNum)
assert.Equal(t, collectionName, task.result.GetCollectionName())
assert.Equal(t, 2, len(task.result.Schema.Fields))
}
func TestDescribeCollectionTask_FilterNamespaceField(t *testing.T) {
mix := NewMixCoordMock()
ctx := context.Background()
InitMetaCache(ctx, mix)
dbName := ""
collectionName := "TestDescribeCollectionTask_FilterNamespaceField"
schema := &schemapb.CollectionSchema{
Name: collectionName,
EnableDynamicField: true,
Fields: []*schemapb.FieldSchema{
{
FieldID: common.StartOfUserFieldID,
Name: "id",
IsPrimaryKey: true,
AutoID: true,
DataType: schemapb.DataType_Int64,
},
{
FieldID: common.StartOfUserFieldID + 1,
Name: "fvec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(8),
},
},
},
{
FieldID: common.StartOfUserFieldID + 2,
Name: common.NamespaceFieldName,
IsPartitionKey: true,
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxLengthKey, Value: "128"}},
IsClusteringKey: false,
},
{
FieldID: common.StartOfUserFieldID + 3,
Name: common.MetaFieldName,
DataType: schemapb.DataType_JSON,
IsDynamic: true,
},
},
}
mix.SetDescribeCollectionFunc(func(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
Schema: schema,
CollectionID: 1,
}, nil
})
task := &describeCollectionTask{
Condition: NewTaskCondition(ctx),
DescribeCollectionRequest: &milvuspb.DescribeCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DescribeCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
},
ctx: ctx,
mixCoord: mix,
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.GetStatus().GetErrorCode())
fieldNames := make(map[string]struct{})
for _, f := range task.result.Schema.Fields {
fieldNames[f.GetName()] = struct{}{}
}
_, hasNamespace := fieldNames[common.NamespaceFieldName]
assert.False(t, hasNamespace)
_, hasMeta := fieldNames[common.MetaFieldName]
assert.False(t, hasMeta)
assert.Equal(t, 2, len(task.result.Schema.Fields))
assert.Equal(t, collectionName, task.result.GetCollectionName())
}
// Security regression: DescribeCollection must redact credentials embedded in
// ExternalSpec.extfs (access_key_id / access_key_value / ssl_ca_cert) before
// returning to the client. Any caller with Describe privilege would otherwise
// be able to read another tenant's object-storage credentials out of the
// persisted collection metadata.
func TestDescribeCollectionTask_RedactsExternalSpecCredentials(t *testing.T) {
mix := NewMixCoordMock()
ctx := context.Background()
InitMetaCache(ctx, mix)
collectionName := "TestDescribeCollectionTask_RedactsExternalSpecCredentials"
rawSpec := `{"format":"parquet","extfs":{"access_key_id":"AKIAEXAMPLE","access_key_value":"SUPERSECRET","region":"us-east-1"}}`
schema := &schemapb.CollectionSchema{
Name: collectionName,
ExternalSource: "s3://my-bucket/data/",
ExternalSpec: rawSpec,
Fields: []*schemapb.FieldSchema{
{
FieldID: common.StartOfUserFieldID,
Name: "id",
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
},
},
}
mix.SetDescribeCollectionFunc(func(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
Schema: schema,
CollectionID: 1,
}, nil
})
task := &describeCollectionTask{
Condition: NewTaskCondition(ctx),
DescribeCollectionRequest: &milvuspb.DescribeCollectionRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_DescribeCollection},
CollectionName: collectionName,
},
ctx: ctx,
mixCoord: mix,
}
assert.NoError(t, task.PreExecute(ctx))
assert.NoError(t, task.Execute(ctx))
// describeCollectionTask now passes ExternalSpec through unredacted —
// redaction lives at the public proxy.DescribeCollection edge so both
// cached and remote provider paths converge through one sanitizer.
// The task-level result must still carry raw creds so that the
// post-task wrapper can choose to redact (user-facing) or not
// (internal callers, if added later, that need raw spec for FFI).
out := task.result.Schema.ExternalSpec
assert.Contains(t, out, "AKIAEXAMPLE",
"task-level result must carry raw spec; redaction is applied at proxy edge")
assert.Contains(t, out, "SUPERSECRET",
"task-level result must carry raw spec; redaction is applied at proxy edge")
}
func TestDescribeCollectionTask_ShardsNum2(t *testing.T) {
mix := NewMixCoordMock()
ctx := context.Background()
InitMetaCache(ctx, mix)
prefix := "TestDescribeCollectionTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
int64Field := "int64"
floatVecField := "fvec"
dim := 128
schema := constructCollectionSchema(int64Field, floatVecField, dim, collectionName)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DropCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
}
mix.CreateCollection(ctx, createColReq)
globalMetaCache.GetCollectionID(ctx, GetCurDBNameFromContextOrDefault(ctx), collectionName)
// CreateCollection
task := &describeCollectionTask{
Condition: NewTaskCondition(ctx),
DescribeCollectionRequest: &milvuspb.DescribeCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DescribeCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
},
ctx: ctx,
mixCoord: mix,
result: nil,
}
task.PreExecute(ctx)
// missing collectionID in globalMetaCache
err = task.Execute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.GetStatus().GetErrorCode())
assert.Equal(t, common.DefaultShardsNum, task.result.ShardsNum)
assert.Equal(t, collectionName, task.result.GetCollectionName())
}
func TestCreatePartitionTask(t *testing.T) {
rc := mocks.NewMockMixCoordClient(t)
mockCache := NewMockCache(t)
mockCache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(newSchemaInfo(&schemapb.CollectionSchema{
EnableDynamicField: false,
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "ID", DataType: schemapb.DataType_Int64},
{FieldID: 101, Name: "Vector", DataType: schemapb.DataType_FloatVector},
},
}), nil)
globalMetaCache = mockCache
ctx := context.Background()
prefix := "TestCreatePartitionTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
partitionName := prefix + funcutil.GenRandomStr()
task := &createPartitionTask{
Condition: NewTaskCondition(ctx),
CreatePartitionRequest: &milvuspb.CreatePartitionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_CreatePartition,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: partitionName,
},
ctx: ctx,
mixCoord: rc,
result: nil,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_CreatePartition, task.Type())
assert.Equal(t, UniqueID(100), task.ID())
assert.Equal(t, Timestamp(100), task.BeginTs())
assert.Equal(t, Timestamp(100), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.GetBase().GetSourceID())
// setup global meta cache
mockCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(100, nil).Once()
rc.EXPECT().CreatePartitionV2(mock.Anything, mock.Anything).Return(&rootcoordpb.CreatePartitionResponse{
Status: merr.Success(),
PartitionID: 1000,
}, nil).Once()
rc.EXPECT().SyncNewCreatedPartition(mock.Anything, mock.Anything).Return(merr.Success(), nil).Once()
err := task.Execute(ctx)
assert.NoError(t, err)
task.CollectionName = "#0xc0de"
err = task.PreExecute(ctx)
assert.Error(t, err)
task.CollectionName = collectionName
task.PartitionName = "#0xc0de"
err = task.PreExecute(ctx)
assert.Error(t, err)
}
func TestDropPartitionTask(t *testing.T) {
ctx := context.Background()
prefix := "TestDropPartitionTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
partitionName := prefix + funcutil.GenRandomStr()
mixc := NewMixCoordMock()
mockCache := NewMockCache(t)
mockCache.On("GetCollectionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(1), nil)
mockCache.On("GetPartitionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(1), nil)
mockCache.On("GetCollectionSchema",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(newSchemaInfo(&schemapb.CollectionSchema{}), nil)
globalMetaCache = mockCache
task := &dropPartitionTask{
Condition: NewTaskCondition(ctx),
DropPartitionRequest: &milvuspb.DropPartitionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DropPartition,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: partitionName,
},
ctx: ctx,
mixCoord: mixc,
result: nil,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_DropPartition, task.Type())
assert.Equal(t, UniqueID(100), task.ID())
assert.Equal(t, Timestamp(100), task.BeginTs())
assert.Equal(t, Timestamp(100), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.GetBase().GetSourceID())
err := task.Execute(ctx)
assert.Error(t, err)
task.CollectionName = "#0xc0de"
err = task.PreExecute(ctx)
assert.Error(t, err)
task.CollectionName = collectionName
task.PartitionName = "#0xc0de"
err = task.PreExecute(ctx)
assert.Error(t, err)
t.Run("get collectionID error", func(t *testing.T) {
mockCache := NewMockCache(t)
mockCache.On("GetCollectionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(1), errors.New("error"))
mockCache.On("GetCollectionSchema",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(newSchemaInfo(&schemapb.CollectionSchema{}), nil)
globalMetaCache = mockCache
task.PartitionName = "partition1"
err = task.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("partition not exist", func(t *testing.T) {
task.PartitionName = "partition2"
mockCache := NewMockCache(t)
mockCache.On("GetPartitionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(0), merr.WrapErrPartitionNotFound(partitionName))
mockCache.On("GetCollectionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(1), nil)
mockCache.On("GetCollectionSchema",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(newSchemaInfo(&schemapb.CollectionSchema{}), nil)
globalMetaCache = mockCache
err = task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("get partition error", func(t *testing.T) {
task.PartitionName = "partition3"
mockCache := NewMockCache(t)
mockCache.On("GetPartitionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(0), errors.New("error"))
mockCache.On("GetCollectionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(1), nil)
mockCache.On("GetCollectionSchema",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(newSchemaInfo(&schemapb.CollectionSchema{}), nil)
globalMetaCache = mockCache
err = task.PreExecute(ctx)
assert.Error(t, err)
})
}
func TestHasPartitionTask(t *testing.T) {
rc := NewMixCoordMock()
defer rc.Close()
ctx := context.Background()
prefix := "TestHasPartitionTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
partitionName := prefix + funcutil.GenRandomStr()
task := &hasPartitionTask{
Condition: NewTaskCondition(ctx),
HasPartitionRequest: &milvuspb.HasPartitionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_HasPartition,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: partitionName,
},
ctx: ctx,
mixCoord: rc,
result: nil,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_HasPartition, task.Type())
assert.Equal(t, UniqueID(100), task.ID())
assert.Equal(t, Timestamp(100), task.BeginTs())
assert.Equal(t, Timestamp(100), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.GetBase().GetSourceID())
err := task.Execute(ctx)
assert.Error(t, err)
task.CollectionName = "#0xc0de"
err = task.PreExecute(ctx)
assert.Error(t, err)
task.CollectionName = collectionName
task.PartitionName = "#0xc0de"
err = task.PreExecute(ctx)
assert.Error(t, err)
}
func TestShowPartitionsTask(t *testing.T) {
mixc := NewMixCoordMock()
defer mixc.Close()
ctx := context.Background()
prefix := "TestShowPartitionsTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
partitionName := prefix + funcutil.GenRandomStr()
task := &showPartitionsTask{
Condition: NewTaskCondition(ctx),
ShowPartitionsRequest: &milvuspb.ShowPartitionsRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_ShowPartitions,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
PartitionNames: []string{partitionName},
Type: milvuspb.ShowType_All,
},
ctx: ctx,
mixCoord: mixc,
result: nil,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_ShowPartitions, task.Type())
assert.Equal(t, UniqueID(100), task.ID())
assert.Equal(t, Timestamp(100), task.BeginTs())
assert.Equal(t, Timestamp(100), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.GetBase().GetSourceID())
err := task.Execute(ctx)
assert.Error(t, err)
task.CollectionName = "#0xc0de"
err = task.PreExecute(ctx)
assert.Error(t, err)
task.CollectionName = collectionName
task.ShowPartitionsRequest.Type = milvuspb.ShowType_InMemory
task.PartitionNames = []string{"#0xc0de"}
err = task.PreExecute(ctx)
assert.Error(t, err)
task.CollectionName = collectionName
task.PartitionNames = []string{partitionName}
task.ShowPartitionsRequest.Type = milvuspb.ShowType_InMemory
err = task.Execute(ctx)
assert.Error(t, err)
}
func TestTask_Int64PrimaryKey(t *testing.T) {
var err error
qc := NewMixCoordMock()
ctx := context.Background()
err = InitMetaCache(ctx, qc)
assert.NoError(t, err)
shardsNum := int32(2)
prefix := "TestTask_int64pk"
dbName := "int64PK"
collectionName := prefix + funcutil.GenRandomStr()
partitionName := prefix + funcutil.GenRandomStr()
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
}
nb := 10
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Types, testInt64Field, false)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createColT := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
assert.NoError(t, createColT.OnEnqueue())
assert.NoError(t, createColT.PreExecute(ctx))
assert.NoError(t, createColT.Execute(ctx))
assert.NoError(t, createColT.PostExecute(ctx))
_, _ = qc.CreatePartition(ctx, &milvuspb.CreatePartitionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_CreatePartition,
MsgID: 0,
Timestamp: 0,
SourceID: paramtable.GetNodeID(),
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: partitionName,
})
collectionID, err := globalMetaCache.GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
dmlChannelsFunc := getDmlChannelsFunc(ctx, qc)
chMgr := newChannelsMgrImpl(dmlChannelsFunc, nil)
pchans, err := chMgr.getChannels(collectionID)
assert.NoError(t, err)
interval := time.Millisecond * 10
tso := newMockTsoAllocator()
ticker := newChannelsTimeTicker(ctx, interval, []string{}, newGetStatisticsFunc(pchans), tso)
_ = ticker.start()
defer ticker.close()
idAllocator, err := allocator.NewIDAllocator(ctx, qc, paramtable.GetNodeID())
assert.NoError(t, err)
_ = idAllocator.Start()
defer idAllocator.Close()
t.Run("insert", func(t *testing.T) {
hash := testutils.GenerateHashKeys(nb)
task := &insertTask{
insertMsg: &BaseInsertTask{
BaseMsg: msgstream.BaseMsg{
HashValues: hash,
},
InsertRequest: &msgpb.InsertRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_Insert,
MsgID: 0,
SourceID: paramtable.GetNodeID(),
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: partitionName,
NumRows: uint64(nb),
Version: msgpb.InsertDataVersion_ColumnBased,
},
},
Condition: NewTaskCondition(ctx),
ctx: ctx,
result: &milvuspb.MutationResult{
Status: merr.Success(),
IDs: nil,
SuccIndex: nil,
ErrIndex: nil,
Acknowledged: false,
InsertCnt: 0,
DeleteCnt: 0,
UpsertCnt: 0,
Timestamp: 0,
},
idAllocator: idAllocator,
chMgr: chMgr,
vChannels: nil,
pChannels: nil,
schema: nil,
}
for fieldName, dataType := range fieldName2Types {
task.insertMsg.FieldsData = append(task.insertMsg.FieldsData, generateFieldData(dataType, fieldName, nb))
}
assert.NoError(t, task.OnEnqueue())
assert.NoError(t, task.PreExecute(ctx))
assert.NoError(t, task.Execute(ctx))
assert.NoError(t, task.PostExecute(ctx))
})
t.Run("simple delete", func(t *testing.T) {
task := &deleteTask{
Condition: NewTaskCondition(ctx),
req: &milvuspb.DeleteRequest{
CollectionName: collectionName,
PartitionName: partitionName,
Expr: "int64 in [0, 1]",
},
idAllocator: idAllocator,
ctx: ctx,
primaryKeys: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{0, 1}}},
},
chMgr: chMgr,
chTicker: ticker,
collectionID: collectionID,
vChannels: []string{"test-ch"},
}
assert.NoError(t, task.OnEnqueue())
assert.NotNil(t, task.TraceCtx())
id := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
task.SetID(id)
assert.Equal(t, id, task.ID())
assert.Equal(t, commonpb.MsgType_Delete, task.Type())
ts := Timestamp(time.Now().UnixNano())
task.SetTs(ts)
assert.Equal(t, ts, task.BeginTs())
assert.Equal(t, ts, task.EndTs())
assert.NoError(t, task.PreExecute(ctx))
assert.NoError(t, task.Execute(ctx))
assert.NoError(t, task.PostExecute(ctx))
})
}
func TestIndexType(t *testing.T) {
rc := NewMixCoordMock()
defer rc.Close()
ctx := context.Background()
shardsNum := int32(2)
prefix := "TestTask_all"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
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,
}
t.Run("invalid type param", func(t *testing.T) {
paramtable.Init()
Params.Save(Params.AutoIndexConfig.Enable.Key, "true")
defer Params.Reset(Params.AutoIndexConfig.Enable.Key)
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Types, testInt64Field, false)
for _, field := range schema.Fields {
dataType := field.GetDataType()
if typeutil.IsVectorType(dataType) {
field.IndexParams = append(field.IndexParams, &commonpb.KeyValuePair{
Key: common.MmapEnabledKey,
Value: "true",
})
break
}
}
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createColT := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: rc,
}
assert.NoError(t, createColT.OnEnqueue())
assert.Error(t, createColT.PreExecute(ctx))
})
}
func TestTask_VarCharPrimaryKey(t *testing.T) {
var err error
mixc := NewMixCoordMock()
ctx := context.Background()
err = InitMetaCache(ctx, mixc)
assert.NoError(t, err)
shardsNum := int32(2)
prefix := "TestTask_all"
dbName := "testvarchar"
collectionName := prefix + funcutil.GenRandomStr()
partitionName := prefix + funcutil.GenRandomStr()
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,
testVarCharField: schemapb.DataType_VarChar,
testFloatVecField: schemapb.DataType_FloatVector,
}
if enableMultipleVectorFields {
fieldName2Types[testBinaryVecField] = schemapb.DataType_BinaryVector
}
nb := 10
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Types, testVarCharField, false)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createColT := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: mixc,
result: nil,
schema: nil,
}
assert.NoError(t, createColT.OnEnqueue())
assert.NoError(t, createColT.PreExecute(ctx))
assert.NoError(t, createColT.Execute(ctx))
assert.NoError(t, createColT.PostExecute(ctx))
_, _ = mixc.CreatePartition(ctx, &milvuspb.CreatePartitionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_CreatePartition,
MsgID: 0,
Timestamp: 0,
SourceID: paramtable.GetNodeID(),
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: partitionName,
})
collectionID, err := globalMetaCache.GetCollectionID(ctx, dbName, collectionName)
assert.NoError(t, err)
dmlChannelsFunc := getDmlChannelsFunc(ctx, mixc)
chMgr := newChannelsMgrImpl(dmlChannelsFunc, nil)
pchans, err := chMgr.getChannels(collectionID)
assert.NoError(t, err)
interval := time.Millisecond * 10
tso := newMockTsoAllocator()
ticker := newChannelsTimeTicker(ctx, interval, []string{}, newGetStatisticsFunc(pchans), tso)
_ = ticker.start()
defer ticker.close()
idAllocator, err := allocator.NewIDAllocator(ctx, mixc, paramtable.GetNodeID())
assert.NoError(t, err)
_ = idAllocator.Start()
defer idAllocator.Close()
t.Run("insert", func(t *testing.T) {
hash := testutils.GenerateHashKeys(nb)
task := &insertTask{
insertMsg: &BaseInsertTask{
BaseMsg: msgstream.BaseMsg{
HashValues: hash,
},
InsertRequest: &msgpb.InsertRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_Insert,
MsgID: 0,
SourceID: paramtable.GetNodeID(),
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: partitionName,
NumRows: uint64(nb),
Version: msgpb.InsertDataVersion_ColumnBased,
},
},
Condition: NewTaskCondition(ctx),
ctx: ctx,
result: &milvuspb.MutationResult{
Status: merr.Success(),
IDs: nil,
SuccIndex: nil,
ErrIndex: nil,
Acknowledged: false,
InsertCnt: 0,
DeleteCnt: 0,
UpsertCnt: 0,
Timestamp: 0,
},
idAllocator: idAllocator,
chMgr: chMgr,
vChannels: nil,
pChannels: nil,
schema: nil,
}
fieldID := common.StartOfUserFieldID
for fieldName, dataType := range fieldName2Types {
task.insertMsg.FieldsData = append(task.insertMsg.FieldsData, generateFieldData(dataType, fieldName, nb))
fieldID++
}
assert.NoError(t, task.OnEnqueue())
assert.NoError(t, task.PreExecute(ctx))
assert.NoError(t, task.Execute(ctx))
assert.NoError(t, task.PostExecute(ctx))
})
t.Run("upsert", func(t *testing.T) {
hash := testutils.GenerateHashKeys(nb)
task := &upsertTask{
upsertMsg: &msgstream.UpsertMsg{
InsertMsg: &BaseInsertTask{
BaseMsg: msgstream.BaseMsg{
HashValues: hash,
},
InsertRequest: &msgpb.InsertRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_Insert,
MsgID: 0,
SourceID: paramtable.GetNodeID(),
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: partitionName,
NumRows: uint64(nb),
Version: msgpb.InsertDataVersion_ColumnBased,
},
},
DeleteMsg: &msgstream.DeleteMsg{
BaseMsg: msgstream.BaseMsg{
HashValues: hash,
},
DeleteRequest: &msgpb.DeleteRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_Delete,
MsgID: 0,
Timestamp: 0,
SourceID: paramtable.GetNodeID(),
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: partitionName,
},
},
},
Condition: NewTaskCondition(ctx),
req: &milvuspb.UpsertRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_Insert,
MsgID: 0,
SourceID: paramtable.GetNodeID(),
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: partitionName,
HashKeys: hash,
NumRows: uint32(nb),
},
ctx: ctx,
result: &milvuspb.MutationResult{
Status: merr.Success(),
IDs: nil,
SuccIndex: nil,
ErrIndex: nil,
Acknowledged: false,
InsertCnt: 0,
DeleteCnt: 0,
UpsertCnt: 0,
Timestamp: 0,
},
idAllocator: idAllocator,
chMgr: chMgr,
chTicker: ticker,
vChannels: nil,
pChannels: nil,
schema: nil,
}
fieldID := common.StartOfUserFieldID
for fieldName, dataType := range fieldName2Types {
task.req.FieldsData = append(task.req.FieldsData, generateFieldData(dataType, fieldName, nb))
fieldID++
}
assert.NoError(t, task.OnEnqueue())
assert.NoError(t, task.PreExecute(ctx))
assert.NoError(t, task.Execute(ctx))
assert.NoError(t, task.PostExecute(ctx))
})
t.Run("simple delete", func(t *testing.T) {
task := &deleteTask{
Condition: NewTaskCondition(ctx),
req: &milvuspb.DeleteRequest{
CollectionName: collectionName,
PartitionName: partitionName,
Expr: "varChar in [\"milvus\", \"test\"]",
},
idAllocator: idAllocator,
ctx: ctx,
chMgr: chMgr,
chTicker: ticker,
vChannels: []string{"test-channel"},
primaryKeys: &schemapb.IDs{
IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: []string{"milvus", "test"}}},
},
collectionID: collectionID,
}
assert.NoError(t, task.OnEnqueue())
assert.NotNil(t, task.TraceCtx())
id := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
task.SetID(id)
assert.Equal(t, id, task.ID())
assert.Equal(t, commonpb.MsgType_Delete, task.Type())
ts := Timestamp(time.Now().UnixNano())
task.SetTs(ts)
assert.Equal(t, ts, task.BeginTs())
assert.Equal(t, ts, task.EndTs())
assert.NoError(t, task.PreExecute(ctx))
assert.NoError(t, task.Execute(ctx))
assert.NoError(t, task.PostExecute(ctx))
})
}
func Test_createIndexTask_getIndexedFieldAndFunction(t *testing.T) {
collectionName := "test"
fieldName := "test"
cit := &createIndexTask{
req: &milvuspb.CreateIndexRequest{
CollectionName: collectionName,
FieldName: fieldName,
},
}
idField := &schemapb.FieldSchema{
FieldID: 100,
Name: "id",
IsPrimaryKey: false,
DataType: schemapb.DataType_FloatVector,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: "dim",
Value: "128",
},
},
AutoID: false,
}
vectorField := &schemapb.FieldSchema{
FieldID: 101,
Name: fieldName,
IsPrimaryKey: false,
DataType: schemapb.DataType_FloatVector,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: "dim",
Value: "128",
},
},
AutoID: false,
}
t.Run("normal", func(t *testing.T) {
cache := NewMockCache(t)
cache.On("GetCollectionSchema",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(newSchemaInfo(&schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
idField,
vectorField,
},
}), nil)
globalMetaCache = cache
err := cit.getIndexedFieldAndFunction(context.Background())
assert.NoError(t, err)
assert.Equal(t, fieldName, cit.fieldSchema.GetName())
})
t.Run("schema not found", func(t *testing.T) {
cache := NewMockCache(t)
cache.On("GetCollectionSchema",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(nil, errors.New("mock"))
globalMetaCache = cache
err := cit.getIndexedFieldAndFunction(context.Background())
assert.Error(t, err)
})
t.Run("field not found", func(t *testing.T) {
otherField := typeutil.Clone(vectorField)
otherField.Name = otherField.Name + "_other"
cache := NewMockCache(t)
cache.On("GetCollectionSchema",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(newSchemaInfo(&schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
idField,
otherField,
},
}), nil)
globalMetaCache = cache
err := cit.getIndexedFieldAndFunction(context.Background())
assert.Error(t, err)
})
}
func Test_fillDimension(t *testing.T) {
// fillDimension moved to indexparamcheck.FillDimension (shared with the
// add-function-field bound-index path); coverage kept here for the proxy flow.
t.Run("scalar", func(t *testing.T) {
f := &schemapb.FieldSchema{
DataType: schemapb.DataType_Int64,
}
assert.NoError(t, indexparamcheck.FillDimension(f, nil))
})
t.Run("no dim in schema", func(t *testing.T) {
f := &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
}
assert.Error(t, indexparamcheck.FillDimension(f, nil))
})
t.Run("dimension mismatch", func(t *testing.T) {
f := &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
}
assert.Error(t, indexparamcheck.FillDimension(f, map[string]string{common.DimKey: "8"}))
})
t.Run("normal", func(t *testing.T) {
f := &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
}
m := map[string]string{}
assert.NoError(t, indexparamcheck.FillDimension(f, m))
assert.Equal(t, "128", m[common.DimKey])
})
}
func Test_checkTrain(t *testing.T) {
t.Run("normal", func(t *testing.T) {
f := &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
}
m := map[string]string{
common.IndexTypeKey: "IVF_FLAT",
"nlist": "1024",
common.MetricTypeKey: "L2",
}
assert.NoError(t, checkTrain(context.TODO(), f, m))
})
t.Run("scalar", func(t *testing.T) {
f := &schemapb.FieldSchema{
DataType: schemapb.DataType_Int64,
}
m := map[string]string{
common.IndexTypeKey: "scalar",
}
assert.Error(t, checkTrain(context.TODO(), f, m))
})
t.Run("dimension mismatch", func(t *testing.T) {
f := &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
}
m := map[string]string{
common.IndexTypeKey: "IVF_FLAT",
"nlist": "1024",
common.MetricTypeKey: "L2",
common.DimKey: "8",
}
assert.Error(t, checkTrain(context.TODO(), f, m))
})
t.Run("nlist test", func(t *testing.T) {
f := &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
}
m := map[string]string{
common.IndexTypeKey: "IVF_FLAT",
common.MetricTypeKey: "L2",
}
assert.NoError(t, checkTrain(context.TODO(), f, m))
})
}
func Test_createIndexTask_PreExecute(t *testing.T) {
collectionName := "test"
fieldName := "test"
cit := &createIndexTask{
req: &milvuspb.CreateIndexRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_CreateIndex,
},
CollectionName: collectionName,
FieldName: fieldName,
},
}
t.Run("normal", func(t *testing.T) {
cache := NewMockCache(t)
cache.On("GetCollectionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(100), nil)
cache.On("GetCollectionSchema",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(newSchemaInfo(&schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: fieldName,
IsPrimaryKey: false,
DataType: schemapb.DataType_FloatVector,
TypeParams: nil,
IndexParams: []*commonpb.KeyValuePair{
{
Key: "dim",
Value: "128",
},
},
AutoID: false,
},
},
}), nil)
cache.On("GetCollectionInfo",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
mock.AnythingOfType("int64"),
).Return(&collectionInfo{}, nil)
globalMetaCache = cache
cit.req.ExtraParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
{
Key: "nlist",
Value: "1024",
},
{
Key: common.MetricTypeKey,
Value: "L2",
},
}
assert.NoError(t, cit.PreExecute(context.Background()))
})
t.Run("collection not found", func(t *testing.T) {
cache := NewMockCache(t)
cache.On("GetCollectionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(0), errors.New("mock"))
globalMetaCache = cache
assert.Error(t, cit.PreExecute(context.Background()))
})
t.Run("index name length exceed 255", func(t *testing.T) {
cache := NewMockCache(t)
cache.On("GetCollectionID",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
).Return(UniqueID(100), nil)
cache.On("GetCollectionInfo",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
mock.AnythingOfType("string"),
mock.AnythingOfType("int64"),
).Return(&collectionInfo{}, nil)
globalMetaCache = cache
for i := 0; i < 256; i++ {
cit.req.IndexName += "a"
}
err := cit.PreExecute(context.Background())
assert.Error(t, err)
})
t.Run("index name start with number", func(t *testing.T) {
cit.req.IndexName = "12a"
err := cit.PreExecute(context.Background())
assert.Error(t, err)
})
t.Run("index name include special characters", func(t *testing.T) {
cit.req.IndexName = "ac#1"
err := cit.PreExecute(context.Background())
assert.Error(t, err)
})
}
func Test_dropCollectionTask_PreExecute(t *testing.T) {
dct := &dropCollectionTask{DropCollectionRequest: &milvuspb.DropCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: "valid", // invalid
}}
globalMetaCache = &MetaCache{}
mockGetCollectionID := mockey.Mock((*MetaCache).GetCollectionID).To(func(ctx context.Context, dbName, collectionName string) (UniqueID, error) {
return 1, nil
}).Build()
defer mockGetCollectionID.UnPatch()
mockRC := mocks.NewMockMixCoordClient(t)
dct.mixCoord = mockRC
ctx := context.Background()
err := dct.PreExecute(ctx)
assert.NoError(t, err)
}
func Test_dropCollectionTask_Execute(t *testing.T) {
mockRC := mocks.NewMockMixCoordClient(t)
mockRC.On("DropCollection",
mock.Anything, // context.Context
mock.Anything, // *milvuspb.DropCollectionRequest
mock.Anything,
).Return(&commonpb.Status{}, func(ctx context.Context, request *milvuspb.DropCollectionRequest, opts ...grpc.CallOption) error {
switch request.GetCollectionName() {
case "c1":
return errors.New("error mock DropCollection")
case "c2":
return merr.WrapErrCollectionNotFound("mock")
default:
return nil
}
})
ctx := context.Background()
dct := &dropCollectionTask{mixCoord: mockRC, DropCollectionRequest: &milvuspb.DropCollectionRequest{CollectionName: "normal"}}
err := dct.Execute(ctx)
assert.NoError(t, err)
dct.CollectionName = "c1"
err = dct.Execute(ctx)
assert.Error(t, err)
dct.CollectionName = "c2"
err = dct.Execute(ctx)
assert.Error(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, dct.result.GetErrorCode())
}
func Test_dropCollectionTask_PostExecute(t *testing.T) {
dct := &dropCollectionTask{}
assert.NoError(t, dct.PostExecute(context.Background()))
}
func Test_truncateCollectionTask_PreExecute(t *testing.T) {
mix := NewMixCoordMock()
ctx := context.Background()
require.NoError(t, InitMetaCache(ctx, mix))
dbName := ""
regularName := "regular_truncate_" + funcutil.GenRandomStr()
externalName := "ext_truncate_" + funcutil.GenRandomStr()
regularSchema := constructCollectionSchema("pk", "fvec", 4, regularName)
regBytes, err := proto.Marshal(regularSchema)
require.NoError(t, err)
_, err = mix.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_CreateCollection},
DbName: dbName, CollectionName: regularName, Schema: regBytes, ShardsNum: 1,
})
require.NoError(t, err)
extSchema := &schemapb.CollectionSchema{
Name: externalName,
ExternalSource: "s3://bucket/path/",
ExternalSpec: `{"format":"parquet"}`,
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true, ExternalField: "id"},
{
FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, ExternalField: "vec",
TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}},
},
},
}
extBytes, err := proto.Marshal(extSchema)
require.NoError(t, err)
_, err = mix.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_CreateCollection},
DbName: dbName, CollectionName: externalName, Schema: extBytes, ShardsNum: 1,
})
require.NoError(t, err)
t.Run("regular collection passes", func(t *testing.T) {
tct := &truncateCollectionTask{TruncateCollectionRequest: &milvuspb.TruncateCollectionRequest{
Base: &commonpb.MsgBase{}, DbName: dbName, CollectionName: regularName,
}}
assert.NoError(t, tct.PreExecute(ctx))
})
t.Run("invalid collection name rejected", func(t *testing.T) {
tct := &truncateCollectionTask{TruncateCollectionRequest: &milvuspb.TruncateCollectionRequest{
Base: &commonpb.MsgBase{}, CollectionName: "#0xc0de",
}}
assert.Error(t, tct.PreExecute(ctx))
})
// Truncate is destructive on internal segments and meaningless for
// external collections (data lives in user object store). Reject up
// front to prevent silent no-op or inconsistent meta state.
t.Run("external collection rejected", func(t *testing.T) {
tct := &truncateCollectionTask{TruncateCollectionRequest: &milvuspb.TruncateCollectionRequest{
Base: &commonpb.MsgBase{}, DbName: dbName, CollectionName: externalName,
}}
err := tct.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "external collection")
assert.Contains(t, err.Error(), "RefreshExternalCollection")
})
}
func Test_truncateCollectionTask_Execute(t *testing.T) {
mockRC := mocks.NewMockMixCoordClient(t)
mockRC.On("TruncateCollection",
mock.Anything, // context.Context
mock.Anything, // *milvuspb.TruncateCollectionRequest
mock.Anything,
).Return(&milvuspb.TruncateCollectionResponse{
Status: merr.Success(),
}, func(ctx context.Context, request *milvuspb.TruncateCollectionRequest, opts ...grpc.CallOption) error {
switch request.GetCollectionName() {
case "c1":
return errors.New("error mock TruncateCollection")
case "c2":
return merr.WrapErrCollectionNotFound("mock")
default:
return nil
}
})
ctx := context.Background()
tct := &truncateCollectionTask{mixCoord: mockRC, TruncateCollectionRequest: &milvuspb.TruncateCollectionRequest{CollectionName: "normal"}}
err := tct.Execute(ctx)
assert.NoError(t, err)
tct.CollectionName = "c1"
err = tct.Execute(ctx)
assert.Error(t, err)
tct.CollectionName = "c2"
err = tct.Execute(ctx)
assert.Error(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, tct.result.GetStatus().GetErrorCode())
}
func Test_truncateCollectionTask_PostExecute(t *testing.T) {
tct := &truncateCollectionTask{}
assert.NoError(t, tct.PostExecute(context.Background()))
}
func Test_loadCollectionTask_Execute(t *testing.T) {
rc := NewMixCoordMock()
dbName := funcutil.GenRandomStr()
collectionName := funcutil.GenRandomStr()
collectionID := UniqueID(1)
// fieldName := funcutil.GenRandomStr()
indexName := funcutil.GenRandomStr()
ctx := context.Background()
indexID := int64(1000)
// failed to get collection id.
_ = InitMetaCache(ctx, rc)
rc.DescribeCollectionFunc = func(ctx context.Context, request *milvuspb.DescribeCollectionRequest, opts ...grpc.CallOption) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
Schema: newTestSchema(),
CollectionID: collectionID,
CollectionName: request.CollectionName,
}, nil
}
lct := &loadCollectionTask{
LoadCollectionRequest: &milvuspb.LoadCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_LoadCollection,
MsgID: 1,
Timestamp: 1,
SourceID: 1,
TargetID: 1,
},
DbName: dbName,
CollectionName: collectionName,
ReplicaNumber: 1,
},
ctx: ctx,
mixCoord: rc,
result: nil,
collectionID: 0,
}
t.Run("indexcoord describe index error", func(t *testing.T) {
err := lct.Execute(ctx)
assert.Error(t, err)
})
t.Run("indexcoord describe index not success", func(t *testing.T) {
rc.DescribeIndexFunc = func(ctx context.Context, request *indexpb.DescribeIndexRequest, opts ...grpc.CallOption) (*indexpb.DescribeIndexResponse, error) {
return &indexpb.DescribeIndexResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "fail reason",
},
}, nil
}
err := lct.Execute(ctx)
assert.Error(t, err)
})
t.Run("no vector index", func(t *testing.T) {
rc.DescribeIndexFunc = func(ctx context.Context, request *indexpb.DescribeIndexRequest, opts ...grpc.CallOption) (*indexpb.DescribeIndexResponse, error) {
return &indexpb.DescribeIndexResponse{
Status: merr.Success(),
IndexInfos: []*indexpb.IndexInfo{
{
CollectionID: collectionID,
FieldID: 100,
IndexName: indexName,
IndexID: indexID,
TypeParams: nil,
IndexParams: nil,
IndexedRows: 1025,
TotalRows: 1025,
State: commonpb.IndexState_Finished,
IndexStateFailReason: "",
IsAutoIndex: false,
UserIndexParams: nil,
},
},
}, nil
}
err := lct.Execute(ctx)
assert.Error(t, err)
})
t.Run("not all vector fields with index", func(t *testing.T) {
vecFields := make([]*schemapb.FieldSchema, 0)
for _, field := range newTestSchema().GetFields() {
if typeutil.IsVectorType(field.GetDataType()) {
vecFields = append(vecFields, field)
}
}
assert.GreaterOrEqual(t, len(vecFields), 2)
rc.DescribeIndexFunc = func(ctx context.Context, request *indexpb.DescribeIndexRequest, opts ...grpc.CallOption) (*indexpb.DescribeIndexResponse, error) {
return &indexpb.DescribeIndexResponse{
Status: merr.Success(),
IndexInfos: []*indexpb.IndexInfo{
{
CollectionID: collectionID,
FieldID: vecFields[0].FieldID,
IndexName: indexName,
IndexID: indexID,
TypeParams: nil,
IndexParams: nil,
IndexedRows: 1025,
TotalRows: 1025,
State: commonpb.IndexState_Finished,
IndexStateFailReason: "",
IsAutoIndex: false,
UserIndexParams: nil,
},
},
}, nil
}
err := lct.Execute(ctx)
assert.Error(t, err)
})
}
func TestLoadCollectionTaskExecuteTextRequiresStorageV3(t *testing.T) {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "false")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
})
oldCache := globalMetaCache
t.Cleanup(func() {
globalMetaCache = oldCache
})
const (
dbName = "db"
collectionName = "text_collection"
collectionID = int64(100)
)
schema := newSchemaInfo(newTextSchemaForStorageV3Test(collectionName))
cache := NewMockCache(t)
cache.EXPECT().GetCollectionID(mock.Anything, dbName, collectionName).Return(collectionID, nil)
cache.EXPECT().GetCollectionSchema(mock.Anything, dbName, collectionName).Return(schema, nil)
globalMetaCache = cache
task := &loadCollectionTask{
LoadCollectionRequest: &milvuspb.LoadCollectionRequest{
Base: commonpbutil.NewMsgBase(),
DbName: dbName,
CollectionName: collectionName,
},
ctx: context.Background(),
}
err := task.Execute(context.Background())
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "TEXT field requires StorageV3")
}
func Test_loadPartitionTask_Execute(t *testing.T) {
qc := NewMixCoordMock()
dbName := funcutil.GenRandomStr()
collectionName := funcutil.GenRandomStr()
collectionID := UniqueID(1)
// fieldName := funcutil.GenRandomStr()
indexName := funcutil.GenRandomStr()
ctx := context.Background()
indexID := int64(1000)
// failed to get collection id.
_ = InitMetaCache(ctx, qc)
qc.DescribeCollectionFunc = func(ctx context.Context, request *milvuspb.DescribeCollectionRequest, opts ...grpc.CallOption) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Status: merr.Success(),
Schema: newTestSchema(),
CollectionID: collectionID,
CollectionName: request.CollectionName,
}, nil
}
lpt := &loadPartitionsTask{
LoadPartitionsRequest: &milvuspb.LoadPartitionsRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_LoadCollection,
MsgID: 1,
Timestamp: 1,
SourceID: 1,
TargetID: 1,
},
DbName: dbName,
CollectionName: collectionName,
ReplicaNumber: 1,
},
ctx: ctx,
mixCoord: qc,
result: nil,
collectionID: 0,
}
t.Run("indexcoord describe index error", func(t *testing.T) {
err := lpt.Execute(ctx)
assert.Error(t, err)
})
t.Run("indexcoord describe index not success", func(t *testing.T) {
qc.DescribeIndexFunc = func(ctx context.Context, request *indexpb.DescribeIndexRequest, opts ...grpc.CallOption) (*indexpb.DescribeIndexResponse, error) {
return &indexpb.DescribeIndexResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "fail reason",
},
}, nil
}
err := lpt.Execute(ctx)
assert.Error(t, err)
})
t.Run("no vector index", func(t *testing.T) {
qc.DescribeIndexFunc = func(ctx context.Context, request *indexpb.DescribeIndexRequest, opts ...grpc.CallOption) (*indexpb.DescribeIndexResponse, error) {
return &indexpb.DescribeIndexResponse{
Status: merr.Success(),
IndexInfos: []*indexpb.IndexInfo{
{
CollectionID: collectionID,
FieldID: 100,
IndexName: indexName,
IndexID: indexID,
TypeParams: nil,
IndexParams: nil,
IndexedRows: 1025,
TotalRows: 1025,
State: commonpb.IndexState_Finished,
IndexStateFailReason: "",
IsAutoIndex: false,
UserIndexParams: nil,
},
},
}, nil
}
err := lpt.Execute(ctx)
assert.Error(t, err)
})
}
func TestCreateResourceGroupTask(t *testing.T) {
mixc := NewMixCoordMock()
defer mixc.Close()
ctx := context.Background()
InitMetaCache(ctx, mixc)
createRGReq := &milvuspb.CreateResourceGroupRequest{
Base: &commonpb.MsgBase{
MsgID: 1,
Timestamp: 2,
TargetID: 3,
},
ResourceGroup: "rg",
}
task := &CreateResourceGroupTask{
CreateResourceGroupRequest: createRGReq,
ctx: ctx,
mixCoord: mixc,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_CreateResourceGroup, task.Type())
assert.Equal(t, UniqueID(1), task.ID())
assert.Equal(t, Timestamp(2), task.BeginTs())
assert.Equal(t, Timestamp(2), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.Base.GetSourceID())
assert.Equal(t, UniqueID(3), task.Base.GetTargetID())
err := task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.ErrorCode)
}
func TestDropResourceGroupTask(t *testing.T) {
mixc := NewMixCoordMock()
defer mixc.Close()
ctx := context.Background()
InitMetaCache(ctx, mixc)
dropRGReq := &milvuspb.DropResourceGroupRequest{
Base: &commonpb.MsgBase{
MsgID: 1,
Timestamp: 2,
TargetID: 3,
},
ResourceGroup: "rg",
}
task := &DropResourceGroupTask{
DropResourceGroupRequest: dropRGReq,
ctx: ctx,
mixCoord: mixc,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_DropResourceGroup, task.Type())
assert.Equal(t, UniqueID(1), task.ID())
assert.Equal(t, Timestamp(2), task.BeginTs())
assert.Equal(t, Timestamp(2), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.Base.GetSourceID())
assert.Equal(t, UniqueID(3), task.Base.GetTargetID())
err := task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.ErrorCode)
}
func TestTransferNodeTask(t *testing.T) {
mixc := NewMixCoordMock()
defer mixc.Close()
ctx := context.Background()
InitMetaCache(ctx, mixc)
req := &milvuspb.TransferNodeRequest{
Base: &commonpb.MsgBase{
MsgID: 1,
Timestamp: 2,
TargetID: 3,
},
SourceResourceGroup: "rg1",
TargetResourceGroup: "rg2",
NumNode: 1,
}
task := &TransferNodeTask{
TransferNodeRequest: req,
ctx: ctx,
mixCoord: mixc,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_TransferNode, task.Type())
assert.Equal(t, UniqueID(1), task.ID())
assert.Equal(t, Timestamp(2), task.BeginTs())
assert.Equal(t, Timestamp(2), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.Base.GetSourceID())
assert.Equal(t, UniqueID(3), task.Base.GetTargetID())
err := task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.ErrorCode)
}
func TestTransferReplicaTask(t *testing.T) {
rc := &MockMixCoordClientInterface{}
ctx := context.Background()
InitMetaCache(ctx, rc)
// make it avoid remote call on rc
globalMetaCache.GetCollectionSchema(context.Background(), GetCurDBNameFromContextOrDefault(ctx), "collection1")
req := &milvuspb.TransferReplicaRequest{
Base: &commonpb.MsgBase{
MsgID: 1,
Timestamp: 2,
TargetID: 3,
},
CollectionName: "collection1",
SourceResourceGroup: "rg1",
TargetResourceGroup: "rg2",
NumReplica: 1,
}
task := &TransferReplicaTask{
TransferReplicaRequest: req,
ctx: ctx,
mixCoord: rc,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_TransferReplica, task.Type())
assert.Equal(t, UniqueID(1), task.ID())
assert.Equal(t, Timestamp(2), task.BeginTs())
assert.Equal(t, Timestamp(2), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.Base.GetSourceID())
assert.Equal(t, UniqueID(3), task.Base.GetTargetID())
err := task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.ErrorCode)
}
func TestListResourceGroupsTask(t *testing.T) {
rc := &MockMixCoordClientInterface{}
rc.listResourceGroups = func(ctx context.Context, request *milvuspb.ListResourceGroupsRequest) (*milvuspb.ListResourceGroupsResponse, error) {
return &milvuspb.ListResourceGroupsResponse{
Status: merr.Success(),
ResourceGroups: []string{meta.DefaultResourceGroupName, "rg"},
}, nil
}
ctx := context.Background()
InitMetaCache(ctx, rc)
req := &milvuspb.ListResourceGroupsRequest{
Base: &commonpb.MsgBase{
MsgID: 1,
Timestamp: 2,
TargetID: 3,
},
}
task := &ListResourceGroupsTask{
ListResourceGroupsRequest: req,
ctx: ctx,
mixCoord: rc,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_ListResourceGroups, task.Type())
assert.Equal(t, UniqueID(1), task.ID())
assert.Equal(t, Timestamp(2), task.BeginTs())
assert.Equal(t, Timestamp(2), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.Base.GetSourceID())
assert.Equal(t, UniqueID(3), task.Base.GetTargetID())
err := task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.GetStatus().GetErrorCode())
groups := task.result.GetResourceGroups()
assert.Contains(t, groups, meta.DefaultResourceGroupName)
assert.Contains(t, groups, "rg")
}
func TestDescribeResourceGroupTask(t *testing.T) {
rc := &MockMixCoordClientInterface{}
rc.describeResourceGroup = func(ctx context.Context, request *querypb.DescribeResourceGroupRequest) (*querypb.DescribeResourceGroupResponse, error) {
return &querypb.DescribeResourceGroupResponse{
Status: merr.Success(),
ResourceGroup: &querypb.ResourceGroupInfo{
Name: "rg",
Capacity: 2,
NumAvailableNode: 1,
NumOutgoingNode: map[int64]int32{1: 1},
NumIncomingNode: map[int64]int32{2: 2},
},
}, nil
}
ctx := context.Background()
InitMetaCache(ctx, rc)
// make it avoid remote call on rc
globalMetaCache.GetCollectionSchema(context.Background(), GetCurDBNameFromContextOrDefault(ctx), "collection1")
globalMetaCache.GetCollectionSchema(context.Background(), GetCurDBNameFromContextOrDefault(ctx), "collection2")
req := &milvuspb.DescribeResourceGroupRequest{
Base: &commonpb.MsgBase{
MsgID: 1,
Timestamp: 2,
TargetID: 3,
},
ResourceGroup: "rg",
}
task := &DescribeResourceGroupTask{
DescribeResourceGroupRequest: req,
ctx: ctx,
mixCoord: rc,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_DescribeResourceGroup, task.Type())
assert.Equal(t, UniqueID(1), task.ID())
assert.Equal(t, Timestamp(2), task.BeginTs())
assert.Equal(t, Timestamp(2), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.Base.GetSourceID())
assert.Equal(t, UniqueID(3), task.Base.GetTargetID())
err := task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.GetStatus().GetErrorCode())
groupInfo := task.result.GetResourceGroup()
outgoingNodeNum := groupInfo.GetNumOutgoingNode()
incomingNodeNum := groupInfo.GetNumIncomingNode()
assert.NotNil(t, outgoingNodeNum["collection1"])
assert.NotNil(t, incomingNodeNum["collection2"])
}
func TestDescribeResourceGroupTaskFailed(t *testing.T) {
rc := &MockMixCoordClientInterface{}
rc.describeResourceGroup = func(ctx context.Context, request *querypb.DescribeResourceGroupRequest) (*querypb.DescribeResourceGroupResponse, error) {
return &querypb.DescribeResourceGroupResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
}, nil
}
ctx := context.Background()
InitMetaCache(ctx, rc)
// make it avoid remote call on rc
globalMetaCache.GetCollectionSchema(context.Background(), GetCurDBNameFromContextOrDefault(ctx), "collection1")
globalMetaCache.GetCollectionSchema(context.Background(), GetCurDBNameFromContextOrDefault(ctx), "collection2")
req := &milvuspb.DescribeResourceGroupRequest{
Base: &commonpb.MsgBase{
MsgID: 1,
Timestamp: 2,
TargetID: 3,
},
ResourceGroup: "rgggg",
}
task := &DescribeResourceGroupTask{
DescribeResourceGroupRequest: req,
ctx: ctx,
mixCoord: rc,
}
task.OnEnqueue()
task.PreExecute(ctx)
assert.Equal(t, commonpb.MsgType_DescribeResourceGroup, task.Type())
assert.Equal(t, UniqueID(1), task.ID())
assert.Equal(t, Timestamp(2), task.BeginTs())
assert.Equal(t, Timestamp(2), task.EndTs())
assert.Equal(t, paramtable.GetNodeID(), task.Base.GetSourceID())
assert.Equal(t, UniqueID(3), task.Base.GetTargetID())
err := task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, task.result.GetStatus().GetErrorCode())
rc.describeResourceGroup = func(ctx context.Context, request *querypb.DescribeResourceGroupRequest) (*querypb.DescribeResourceGroupResponse, error) {
return &querypb.DescribeResourceGroupResponse{
Status: merr.Success(),
ResourceGroup: &querypb.ResourceGroupInfo{
Name: "rg",
Capacity: 2,
NumAvailableNode: 1,
NumOutgoingNode: map[int64]int32{3: 1},
NumIncomingNode: map[int64]int32{4: 2},
},
}, nil
}
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Len(t, task.result.ResourceGroup.NumOutgoingNode, 0)
assert.Len(t, task.result.ResourceGroup.NumIncomingNode, 0)
}
func TestCreateCollectionTaskWithPartitionKey(t *testing.T) {
rc := NewMixCoordMock()
paramtable.Init()
defer rc.Close()
cache := globalMetaCache
defer func() { globalMetaCache = cache }()
mockCache := NewMockCache(t)
mockCache.EXPECT().GetDatabaseInfo(mock.Anything, mock.Anything).Return(&databaseInfo{}, nil).Maybe()
globalMetaCache = mockCache
ctx := context.Background()
shardsNum := common.DefaultShardsNum
prefix := "TestCreateCollectionTaskWithPartitionKey"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
int64Field := &schemapb.FieldSchema{
Name: "int64",
DataType: schemapb.DataType_Int64,
IsPrimaryKey: true,
}
varCharField := &schemapb.FieldSchema{
Name: "varChar",
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{
Key: "max_length",
Value: strconv.Itoa(testMaxVarCharLength),
},
{
Key: "enable_analyzer",
Value: "true",
},
},
}
floatVecField := &schemapb.FieldSchema{
Name: "fvec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: "dim",
Value: strconv.Itoa(testVecDim),
},
},
}
sparseVecField := &schemapb.FieldSchema{
Name: "sparse",
DataType: schemapb.DataType_SparseFloatVector,
}
partitionKeyField := &schemapb.FieldSchema{
Name: "partition_key",
DataType: schemapb.DataType_Int64,
IsPartitionKey: true,
}
schema := &schemapb.CollectionSchema{
Name: collectionName,
Fields: []*schemapb.FieldSchema{int64Field, varCharField, partitionKeyField, floatVecField},
}
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: rc,
result: nil,
schema: nil,
}
t.Run("PreExecute", func(t *testing.T) {
defer Params.Reset(Params.RootCoordCfg.MaxPartitionNum.Key)
var err error
// test default num partitions
err = task.PreExecute(ctx)
assert.NoError(t, err)
assert.Equal(t, common.DefaultPartitionsWithPartitionKey, task.GetNumPartitions())
Params.Save(Params.RootCoordCfg.MaxPartitionNum.Key, "16")
task.NumPartitions = 0
err = task.PreExecute(ctx)
assert.NoError(t, err)
assert.Equal(t, int64(16), task.GetNumPartitions())
Params.Reset(Params.RootCoordCfg.MaxPartitionNum.Key)
// test specify num partition without partition key field
partitionKeyField.IsPartitionKey = false
task.NumPartitions = common.DefaultPartitionsWithPartitionKey * 2
marshaledSchema, err = proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = marshaledSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
partitionKeyField.IsPartitionKey = true
// test multi partition key field
varCharField.IsPartitionKey = true
marshaledSchema, err = proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = marshaledSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
varCharField.IsPartitionKey = false
// test partitions < 0
task.NumPartitions = -2
marshaledSchema, err = proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = marshaledSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
task.NumPartitions = 1000
// test partition key type not in [int64, varChar]
partitionKeyField.DataType = schemapb.DataType_FloatVector
marshaledSchema, err = proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = marshaledSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
partitionKeyField.DataType = schemapb.DataType_Int64
// test partition key set nullable == true
partitionKeyField.Nullable = true
marshaledSchema, err = proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = marshaledSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
partitionKeyField.DataType = schemapb.DataType_Int64
partitionKeyField.Nullable = false
// test partition key field not primary key field
primaryField, _ := typeutil.GetPrimaryFieldSchema(schema)
primaryField.IsPartitionKey = true
marshaledSchema, err = proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = marshaledSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
primaryField.IsPartitionKey = false
// test partition num too large
Params.Save(Params.RootCoordCfg.MaxPartitionNum.Key, "16")
marshaledSchema, err = proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = marshaledSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
Params.Reset(Params.RootCoordCfg.MaxPartitionNum.Key)
marshaledSchema, err = proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = marshaledSchema
err = task.PreExecute(ctx)
assert.NoError(t, err)
// test schema with function
// invalid function
schema.Functions = []*schemapb.FunctionSchema{
{Name: "test", Type: schemapb.FunctionType_BM25, InputFieldNames: []string{"invalid name"}},
}
marshaledSchema, err = proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = marshaledSchema
err = task.PreExecute(ctx)
assert.Error(t, err)
// normal case
schema.Fields = append(schema.Fields, sparseVecField)
schema.Functions = []*schemapb.FunctionSchema{
{Name: "test", Type: schemapb.FunctionType_BM25, InputFieldNames: []string{varCharField.Name}, OutputFieldNames: []string{sparseVecField.Name}},
}
marshaledSchema, err = proto.Marshal(schema)
assert.NoError(t, err)
task.Schema = marshaledSchema
err = task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("Execute", func(t *testing.T) {
err = task.Execute(ctx)
assert.NoError(t, err)
// check default partitions
err = InitMetaCache(ctx, rc)
assert.NoError(t, err)
partitionNames, err := getDefaultPartitionsInPartitionKeyMode(ctx, "", task.CollectionName)
assert.NoError(t, err)
assert.Equal(t, task.GetNumPartitions(), int64(len(partitionNames)))
createPartitionTask := &createPartitionTask{
Condition: NewTaskCondition(ctx),
CreatePartitionRequest: &milvuspb.CreatePartitionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: "new_partition",
},
ctx: ctx,
mixCoord: rc,
}
err = createPartitionTask.PreExecute(ctx)
assert.Error(t, err)
dropPartitionTask := &dropPartitionTask{
Condition: NewTaskCondition(ctx),
DropPartitionRequest: &milvuspb.DropPartitionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: dbName,
CollectionName: collectionName,
PartitionName: "new_partition",
},
ctx: ctx,
mixCoord: rc,
}
err = dropPartitionTask.PreExecute(ctx)
assert.Error(t, err)
loadPartitionTask := &loadPartitionsTask{
Condition: NewTaskCondition(ctx),
LoadPartitionsRequest: &milvuspb.LoadPartitionsRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: dbName,
CollectionName: collectionName,
PartitionNames: []string{"_default_0"},
},
ctx: ctx,
}
err = loadPartitionTask.PreExecute(ctx)
assert.Error(t, err)
releasePartitionsTask := &releasePartitionsTask{
Condition: NewTaskCondition(ctx),
ReleasePartitionsRequest: &milvuspb.ReleasePartitionsRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: dbName,
CollectionName: collectionName,
PartitionNames: []string{"_default_0"},
},
ctx: ctx,
}
err = releasePartitionsTask.PreExecute(ctx)
assert.Error(t, err)
})
}
func TestPartitionKey(t *testing.T) {
qc := NewMixCoordMock()
defer qc.Close()
ctx := context.Background()
err := InitMetaCache(ctx, qc)
assert.NoError(t, err)
shardsNum := common.DefaultShardsNum
prefix := "TestInsertTaskWithPartitionKey"
collectionName := prefix + funcutil.GenRandomStr()
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type["int64_field"] = schemapb.DataType_Int64
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
fieldName2Type["fvec_field"] = schemapb.DataType_FloatVector
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, "int64_field", false)
partitionKeyField := &schemapb.FieldSchema{
Name: "partition_key_field",
DataType: schemapb.DataType_Int64,
IsPartitionKey: true,
}
fieldName2Type["partition_key_field"] = schemapb.DataType_Int64
schema.Fields = append(schema.Fields, partitionKeyField)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: "",
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
NumPartitions: common.DefaultPartitionsWithPartitionKey,
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
err = createCollectionTask.PreExecute(ctx)
assert.NoError(t, err)
err = createCollectionTask.Execute(ctx)
assert.NoError(t, err)
collectionID, err := globalMetaCache.GetCollectionID(ctx, GetCurDBNameFromContextOrDefault(ctx), collectionName)
assert.NoError(t, err)
dmlChannelsFunc := getDmlChannelsFunc(ctx, qc)
chMgr := newChannelsMgrImpl(dmlChannelsFunc, nil)
pchans, err := chMgr.getChannels(collectionID)
assert.NoError(t, err)
interval := time.Millisecond * 10
tso := newMockTsoAllocator()
ticker := newChannelsTimeTicker(ctx, interval, []string{}, newGetStatisticsFunc(pchans), tso)
_ = ticker.start()
defer ticker.close()
idAllocator, err := allocator.NewIDAllocator(ctx, qc, paramtable.GetNodeID())
assert.NoError(t, err)
_ = idAllocator.Start()
defer idAllocator.Close()
partitionNames, err := getDefaultPartitionsInPartitionKeyMode(ctx, "", collectionName)
assert.NoError(t, err)
assert.Equal(t, common.DefaultPartitionsWithPartitionKey, int64(len(partitionNames)))
nb := 10
fieldID := common.StartOfUserFieldID
fieldDatas := make([]*schemapb.FieldData, 0)
for fieldName, dataType := range fieldName2Type {
fieldData := generateFieldData(dataType, fieldName, nb)
fieldData.FieldId = int64(fieldID)
fieldDatas = append(fieldDatas, generateFieldData(dataType, fieldName, nb))
fieldID++
}
t.Run("Insert", func(t *testing.T) {
it := &insertTask{
insertMsg: &BaseInsertTask{
BaseMsg: msgstream.BaseMsg{},
InsertRequest: &msgpb.InsertRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_Insert,
MsgID: 0,
SourceID: paramtable.GetNodeID(),
},
CollectionName: collectionName,
FieldsData: fieldDatas,
NumRows: uint64(nb),
Version: msgpb.InsertDataVersion_ColumnBased,
},
},
Condition: NewTaskCondition(ctx),
ctx: ctx,
result: &milvuspb.MutationResult{
Status: merr.Success(),
IDs: nil,
SuccIndex: nil,
ErrIndex: nil,
Acknowledged: false,
InsertCnt: 0,
DeleteCnt: 0,
UpsertCnt: 0,
Timestamp: 0,
},
idAllocator: idAllocator,
chMgr: chMgr,
vChannels: nil,
pChannels: nil,
schema: nil,
}
// don't support specify partition name if use partition key
it.insertMsg.PartitionName = partitionNames[0]
assert.Error(t, it.PreExecute(ctx))
it.insertMsg.PartitionName = ""
assert.NoError(t, it.OnEnqueue())
assert.NoError(t, it.PreExecute(ctx))
assert.NoError(t, it.Execute(ctx))
assert.NoError(t, it.PostExecute(ctx))
})
t.Run("Upsert", func(t *testing.T) {
hash := testutils.GenerateHashKeys(nb)
ut := &upsertTask{
ctx: ctx,
Condition: NewTaskCondition(ctx),
baseMsg: msgstream.BaseMsg{
HashValues: hash,
},
req: &milvuspb.UpsertRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_Upsert),
commonpbutil.WithSourceID(paramtable.GetNodeID()),
),
CollectionName: collectionName,
FieldsData: fieldDatas,
NumRows: uint32(nb),
},
result: &milvuspb.MutationResult{
Status: merr.Success(),
IDs: &schemapb.IDs{
IdField: nil,
},
},
idAllocator: idAllocator,
chMgr: chMgr,
chTicker: ticker,
}
// don't support specify partition name if use partition key
ut.req.PartitionName = partitionNames[0]
assert.Error(t, ut.PreExecute(ctx))
ut.req.PartitionName = ""
assert.NoError(t, ut.OnEnqueue())
assert.NoError(t, ut.PreExecute(ctx))
assert.NoError(t, ut.Execute(ctx))
assert.NoError(t, ut.PostExecute(ctx))
})
t.Run("delete", func(t *testing.T) {
dt := &deleteTask{
Condition: NewTaskCondition(ctx),
req: &milvuspb.DeleteRequest{
CollectionName: collectionName,
Expr: "int64_field in [0, 1]",
},
ctx: ctx,
primaryKeys: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{0, 1}}},
},
idAllocator: idAllocator,
chMgr: chMgr,
chTicker: ticker,
collectionID: collectionID,
vChannels: []string{"test-channel"},
}
dt.req.PartitionName = ""
assert.NoError(t, dt.PreExecute(ctx))
assert.NoError(t, dt.Execute(ctx))
assert.NoError(t, dt.PostExecute(ctx))
})
t.Run("search", func(t *testing.T) {
searchTask := &searchTask{
ctx: ctx,
SearchRequest: &internalpb.SearchRequest{
Base: &commonpb.MsgBase{},
},
request: &milvuspb.SearchRequest{
CollectionName: collectionName,
Nq: 1,
},
mixCoord: qc,
tr: timerecord.NewTimeRecorder("test-search"),
}
// don't support specify partition name if use partition key
searchTask.request.PartitionNames = partitionNames
err = searchTask.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("query", func(t *testing.T) {
queryTask := &queryTask{
ctx: ctx,
RetrieveRequest: &internalpb.RetrieveRequest{
QueryLabel: "query",
Base: &commonpb.MsgBase{},
},
request: &milvuspb.QueryRequest{
CollectionName: collectionName,
},
mixCoord: qc,
}
// don't support specify partition name if use partition key
queryTask.request.PartitionNames = partitionNames
err = queryTask.PreExecute(ctx)
assert.Error(t, err)
})
}
func TestDefaultPartition(t *testing.T) {
qc := NewMixCoordMock()
ctx := context.Background()
err := InitMetaCache(ctx, qc)
assert.NoError(t, err)
shardsNum := common.DefaultShardsNum
prefix := "TestInsertTaskWithPartitionKey"
collectionName := prefix + funcutil.GenRandomStr()
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type["int64_field"] = schemapb.DataType_Int64
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
fieldName2Type["fvec_field"] = schemapb.DataType_FloatVector
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, "int64_field", false)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
t.Run("create collection", func(t *testing.T) {
createCollectionTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: "",
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
err = createCollectionTask.PreExecute(ctx)
assert.NoError(t, err)
err = createCollectionTask.Execute(ctx)
assert.NoError(t, err)
})
collectionID, err := globalMetaCache.GetCollectionID(ctx, GetCurDBNameFromContextOrDefault(ctx), collectionName)
assert.NoError(t, err)
dmlChannelsFunc := getDmlChannelsFunc(ctx, qc)
chMgr := newChannelsMgrImpl(dmlChannelsFunc, nil)
pchans, err := chMgr.getChannels(collectionID)
assert.NoError(t, err)
interval := time.Millisecond * 10
tso := newMockTsoAllocator()
ticker := newChannelsTimeTicker(ctx, interval, []string{}, newGetStatisticsFunc(pchans), tso)
_ = ticker.start()
defer ticker.close()
idAllocator, err := allocator.NewIDAllocator(ctx, qc, paramtable.GetNodeID())
assert.NoError(t, err)
_ = idAllocator.Start()
defer idAllocator.Close()
nb := 10
fieldID := common.StartOfUserFieldID
fieldDatas := make([]*schemapb.FieldData, 0)
for fieldName, dataType := range fieldName2Type {
fieldData := generateFieldData(dataType, fieldName, nb)
fieldData.FieldId = int64(fieldID)
fieldDatas = append(fieldDatas, generateFieldData(dataType, fieldName, nb))
fieldID++
}
t.Run("Insert", func(t *testing.T) {
it := &insertTask{
insertMsg: &BaseInsertTask{
BaseMsg: msgstream.BaseMsg{},
InsertRequest: &msgpb.InsertRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_Insert,
MsgID: 0,
SourceID: paramtable.GetNodeID(),
},
CollectionName: collectionName,
FieldsData: fieldDatas,
NumRows: uint64(nb),
Version: msgpb.InsertDataVersion_ColumnBased,
},
},
Condition: NewTaskCondition(ctx),
ctx: ctx,
result: &milvuspb.MutationResult{
Status: merr.Success(),
IDs: nil,
SuccIndex: nil,
ErrIndex: nil,
Acknowledged: false,
InsertCnt: 0,
DeleteCnt: 0,
UpsertCnt: 0,
Timestamp: 0,
},
idAllocator: idAllocator,
chMgr: chMgr,
vChannels: nil,
pChannels: nil,
schema: nil,
}
it.insertMsg.PartitionName = ""
assert.NoError(t, it.OnEnqueue())
assert.NoError(t, it.PreExecute(ctx))
assert.NoError(t, it.Execute(ctx))
assert.NoError(t, it.PostExecute(ctx))
})
t.Run("Upsert", func(t *testing.T) {
hash := testutils.GenerateHashKeys(nb)
ut := &upsertTask{
ctx: ctx,
Condition: NewTaskCondition(ctx),
baseMsg: msgstream.BaseMsg{
HashValues: hash,
},
req: &milvuspb.UpsertRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_Upsert),
commonpbutil.WithSourceID(paramtable.GetNodeID()),
),
CollectionName: collectionName,
FieldsData: fieldDatas,
NumRows: uint32(nb),
},
result: &milvuspb.MutationResult{
Status: merr.Success(),
IDs: &schemapb.IDs{
IdField: nil,
},
},
idAllocator: idAllocator,
chMgr: chMgr,
chTicker: ticker,
}
ut.req.PartitionName = ""
assert.NoError(t, ut.OnEnqueue())
assert.NoError(t, ut.PreExecute(ctx))
assert.NoError(t, ut.Execute(ctx))
assert.NoError(t, ut.PostExecute(ctx))
})
t.Run("delete", func(t *testing.T) {
dt := &deleteTask{
Condition: NewTaskCondition(ctx),
req: &milvuspb.DeleteRequest{
CollectionName: collectionName,
Expr: "int64_field in [0, 1]",
},
ctx: ctx,
primaryKeys: &schemapb.IDs{
IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{0, 1}}},
},
idAllocator: idAllocator,
chMgr: chMgr,
chTicker: ticker,
collectionID: collectionID,
vChannels: []string{"test-channel"},
}
dt.req.PartitionName = ""
assert.NoError(t, dt.PreExecute(ctx))
assert.NoError(t, dt.Execute(ctx))
assert.NoError(t, dt.PostExecute(ctx))
})
}
func TestClusteringKey(t *testing.T) {
qc := NewMixCoordMock()
ctx := context.Background()
err := InitMetaCache(ctx, qc)
assert.NoError(t, err)
shardsNum := common.DefaultShardsNum
prefix := "TestClusteringKey"
collectionName := prefix + funcutil.GenRandomStr()
t.Run("create collection normal", func(t *testing.T) {
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type["int64_field"] = schemapb.DataType_Int64
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, "int64_field", false)
fieldName2Type["cluster_key_field"] = schemapb.DataType_Int64
clusterKeyField := &schemapb.FieldSchema{
Name: "cluster_key_field",
DataType: schemapb.DataType_Int64,
IsClusteringKey: true,
}
schema.Fields = append(schema.Fields, clusterKeyField)
vecField := &schemapb.FieldSchema{
Name: "fvec_field",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(testVecDim),
},
},
}
schema.Fields = append(schema.Fields, vecField)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: "",
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
err = createCollectionTask.PreExecute(ctx)
assert.NoError(t, err)
err = createCollectionTask.Execute(ctx)
assert.NoError(t, err)
})
t.Run("create collection not support more than one clustering key", func(t *testing.T) {
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type["int64_field"] = schemapb.DataType_Int64
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, "int64_field", false)
fieldName2Type["cluster_key_field"] = schemapb.DataType_Int64
clusterKeyField := &schemapb.FieldSchema{
Name: "cluster_key_field",
DataType: schemapb.DataType_Int64,
IsClusteringKey: true,
}
schema.Fields = append(schema.Fields, clusterKeyField)
clusterKeyField2 := &schemapb.FieldSchema{
Name: "cluster_key_field2",
DataType: schemapb.DataType_Int64,
IsClusteringKey: true,
}
schema.Fields = append(schema.Fields, clusterKeyField2)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: "",
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
err = createCollectionTask.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("create collection with vector clustering key", func(t *testing.T) {
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type["int64_field"] = schemapb.DataType_Int64
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, "int64_field", false)
clusterKeyField := &schemapb.FieldSchema{
Name: "vec_field",
DataType: schemapb.DataType_FloatVector,
IsClusteringKey: true,
}
schema.Fields = append(schema.Fields, clusterKeyField)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: "",
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
err = createCollectionTask.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("create collection with json clustering key not supported", func(t *testing.T) {
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type["int64_field"] = schemapb.DataType_Int64
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, "int64_field", false)
jsonField := &schemapb.FieldSchema{
Name: "json_field",
DataType: schemapb.DataType_JSON,
IsClusteringKey: true,
}
schema.Fields = append(schema.Fields, jsonField)
vecField := &schemapb.FieldSchema{
Name: "fvec_field",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(testVecDim),
},
},
}
schema.Fields = append(schema.Fields, vecField)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: "",
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
err = createCollectionTask.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("create collection with bool clustering key not supported", func(t *testing.T) {
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type["int64_field"] = schemapb.DataType_Int64
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, "int64_field", false)
boolField := &schemapb.FieldSchema{
Name: "bool_field",
DataType: schemapb.DataType_Bool,
IsClusteringKey: true,
}
schema.Fields = append(schema.Fields, boolField)
vecField := &schemapb.FieldSchema{
Name: "fvec_field",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(testVecDim),
},
},
}
schema.Fields = append(schema.Fields, vecField)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: "",
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
err = createCollectionTask.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("create collection with array clustering key not supported", func(t *testing.T) {
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type["int64_field"] = schemapb.DataType_Int64
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
schema := constructCollectionSchemaByDataType(collectionName, fieldName2Type, "int64_field", false)
arrayField := &schemapb.FieldSchema{
Name: "array_field",
DataType: schemapb.DataType_Array,
IsClusteringKey: true,
ElementType: schemapb.DataType_Int64,
}
schema.Fields = append(schema.Fields, arrayField)
vecField := &schemapb.FieldSchema{
Name: "fvec_field",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(testVecDim),
},
},
}
schema.Fields = append(schema.Fields, vecField)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: "",
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
err = createCollectionTask.PreExecute(ctx)
assert.Error(t, err)
})
}
func TestAlterCollectionCheckLoaded(t *testing.T) {
qc := NewMixCoordMock()
InitMetaCache(context.Background(), qc)
collectionName := "test_alter_collection_check_loaded"
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DropCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
Schema: nil,
ShardsNum: 1,
}
qc.CreateCollection(context.Background(), createColReq)
resp, err := qc.DescribeCollection(context.Background(), &milvuspb.DescribeCollectionRequest{CollectionName: collectionName})
assert.NoError(t, err)
qc.ShowLoadCollectionsFunc = func(ctx context.Context, req *querypb.ShowCollectionsRequest, opts ...grpc.CallOption) (*querypb.ShowCollectionsResponse, error) {
return &querypb.ShowCollectionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
CollectionIDs: []int64{resp.CollectionID},
InMemoryPercentages: []int64{100},
}, nil
}
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
Properties: []*commonpb.KeyValuePair{{Key: common.MmapEnabledKey, Value: "true"}},
},
mixCoord: qc,
}
err = task.PreExecute(context.Background())
assert.Equal(t, merr.Code(merr.ErrCollectionLoaded), merr.Code(err))
}
func TestAlterCollectionTaskValidateTTLAndTTLField(t *testing.T) {
qc := NewMixCoordMock()
ctx := context.Background()
cache := globalMetaCache
defer func() { globalMetaCache = cache }()
err := InitMetaCache(ctx, qc)
assert.NoError(t, err)
createCollectionWithProps := func(colName string, props []*commonpb.KeyValuePair) {
schema := &schemapb.CollectionSchema{
Name: colName,
Description: "ttl alter test",
AutoID: false,
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", IsPrimaryKey: true, DataType: schemapb.DataType_Int64},
{FieldID: 101, Name: "ttl", DataType: schemapb.DataType_Timestamptz, Nullable: true},
{FieldID: 102, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "8"}}},
},
Properties: props,
}
bs, err := proto.Marshal(schema)
assert.NoError(t, err)
_, err = qc.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_CreateCollection},
DbName: "",
CollectionName: colName,
Schema: bs,
ShardsNum: 1,
Properties: props,
})
assert.NoError(t, err)
}
t.Run("mutual exclusion: existing ttl.seconds, alter ttl.field should fail", func(t *testing.T) {
col := "alter_ttl_seconds_then_field_" + funcutil.GenRandomStr()
createCollectionWithProps(col, []*commonpb.KeyValuePair{{Key: common.CollectionTTLConfigKey, Value: "3600"}})
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
Properties: []*commonpb.KeyValuePair{{Key: common.CollectionTTLFieldKey, Value: "ttl"}},
},
ctx: ctx,
mixCoord: qc,
}
err := task.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("mutual exclusion: existing ttl.field, alter ttl.seconds should fail", func(t *testing.T) {
col := "alter_ttl_field_then_seconds_" + funcutil.GenRandomStr()
createCollectionWithProps(col, []*commonpb.KeyValuePair{{Key: common.CollectionTTLFieldKey, Value: "ttl"}})
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
Properties: []*commonpb.KeyValuePair{{Key: common.CollectionTTLConfigKey, Value: "10"}},
},
ctx: ctx,
mixCoord: qc,
}
err := task.PreExecute(ctx)
assert.Error(t, err)
})
t.Run("ttl.field must exist in schema", func(t *testing.T) {
col := "alter_ttl_field_invalid_" + funcutil.GenRandomStr()
createCollectionWithProps(col, nil)
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
Properties: []*commonpb.KeyValuePair{{Key: common.CollectionTTLFieldKey, Value: "not_exist"}},
},
ctx: ctx,
mixCoord: qc,
}
assert.Error(t, task.PreExecute(ctx))
})
t.Run("ttl.field must be timestamptz", func(t *testing.T) {
col := "alter_ttl_field_invalid_type_" + funcutil.GenRandomStr()
createCollectionWithProps(col, nil)
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
Properties: []*commonpb.KeyValuePair{{Key: common.CollectionTTLFieldKey, Value: "pk"}},
},
ctx: ctx,
mixCoord: qc,
}
assert.Error(t, task.PreExecute(ctx))
})
t.Run("delete ttl.seconds should pass", func(t *testing.T) {
col := "alter_delete_ttl_seconds_" + funcutil.GenRandomStr()
createCollectionWithProps(col, []*commonpb.KeyValuePair{{Key: common.CollectionTTLConfigKey, Value: "3600"}})
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
DeleteKeys: []string{common.CollectionTTLConfigKey},
},
ctx: ctx,
mixCoord: qc,
}
assert.NoError(t, task.PreExecute(ctx))
})
t.Run("delete ttl.field should pass", func(t *testing.T) {
col := "alter_delete_ttl_field_" + funcutil.GenRandomStr()
createCollectionWithProps(col, []*commonpb.KeyValuePair{{Key: common.CollectionTTLFieldKey, Value: "ttl"}})
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: col,
DeleteKeys: []string{common.CollectionTTLFieldKey},
},
ctx: ctx,
mixCoord: qc,
}
assert.NoError(t, task.PreExecute(ctx))
})
}
func TestAlterCollectionTaskValidateDescription(t *testing.T) {
qc := NewMixCoordMock()
ctx := context.Background()
cache := globalMetaCache
defer func() { globalMetaCache = cache }()
err := InitMetaCache(ctx, qc)
assert.NoError(t, err)
maxLen := Params.ProxyCfg.MaxCollectionDescriptionLength.GetAsInt()
oversized := strings.Repeat("a", maxLen+1)
differentOversized := strings.Repeat("b", maxLen+1)
// Create a collection through the mock coordinator directly so we can seed a
// pre-existing oversized description, which the proxy CreateCollection guard
// would normally reject. This simulates a collection created before the
// description length limit existed.
createCollection := func(description string) string {
collectionName := "alter_description_" + funcutil.GenRandomStr()
schema := &schemapb.CollectionSchema{
Name: collectionName,
Description: description,
AutoID: false,
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", IsPrimaryKey: true, DataType: schemapb.DataType_Int64},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "8"}}},
},
}
bs, err := proto.Marshal(schema)
assert.NoError(t, err)
status, err := qc.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_CreateCollection},
DbName: dbName,
CollectionName: collectionName,
Schema: bs,
ShardsNum: 1,
})
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, status.GetErrorCode())
return collectionName
}
runAlter := func(collectionName string, properties []*commonpb.KeyValuePair) error {
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
Properties: properties,
},
ctx: ctx,
mixCoord: qc,
}
return task.PreExecute(ctx)
}
// A collection whose stored description is within the limit.
validCollection := createCollection("short description")
// A collection whose stored description already exceeds the current limit.
legacyCollection := createCollection(oversized)
t.Run("changing a valid description to an oversized one is rejected", func(t *testing.T) {
err := runAlter(validCollection, []*commonpb.KeyValuePair{
{Key: common.CollectionDescription, Value: oversized},
})
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("resending an unchanged oversized description while altering TTL is accepted", func(t *testing.T) {
err := runAlter(legacyCollection, []*commonpb.KeyValuePair{
{Key: common.CollectionDescription, Value: oversized},
{Key: common.CollectionTTLConfigKey, Value: "10"},
})
assert.NoError(t, err)
})
t.Run("changing an oversized description to a different oversized value is rejected", func(t *testing.T) {
err := runAlter(legacyCollection, []*commonpb.KeyValuePair{
{Key: common.CollectionDescription, Value: differentOversized},
})
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("shrinking an oversized description to a valid value is accepted", func(t *testing.T) {
err := runAlter(legacyCollection, []*commonpb.KeyValuePair{
{Key: common.CollectionDescription, Value: "now within the limit"},
})
assert.NoError(t, err)
})
t.Run("with duplicate description keys, a changed oversized value is still rejected", func(t *testing.T) {
// The unchanged value is skipped, but the other (changed) oversized value
// is still validated per-property, so the alter is rejected.
err := runAlter(legacyCollection, []*commonpb.KeyValuePair{
{Key: common.CollectionDescription, Value: oversized},
{Key: common.CollectionDescription, Value: differentOversized},
})
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
}
func mockVectorIndexForCollection(t *testing.T, ctx context.Context, qc *MixCoordMock, colName string) {
t.Helper()
resp, err := qc.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{DbName: dbName, CollectionName: colName})
assert.NoError(t, err)
var vecFieldID int64
for _, field := range resp.Schema.Fields {
if field.DataType == schemapb.DataType_FloatVector {
vecFieldID = field.FieldID
break
}
}
assert.NotEqual(t, int64(0), vecFieldID)
qc.DescribeIndexFunc = func(ctx context.Context, request *indexpb.DescribeIndexRequest, opts ...grpc.CallOption) (*indexpb.DescribeIndexResponse, error) {
return &indexpb.DescribeIndexResponse{
Status: merr.Success(),
IndexInfos: []*indexpb.IndexInfo{
{FieldID: vecFieldID},
},
}, nil
}
t.Cleanup(func() {
qc.DescribeIndexFunc = nil
})
}
func TestTaskPartitionKeyIsolation(t *testing.T) {
qc := NewMixCoordMock()
ctx := context.Background()
err := InitMetaCache(ctx, qc)
assert.NoError(t, err)
shardsNum := common.DefaultShardsNum
prefix := "TestPartitionKeyIsolation"
collectionName := prefix + funcutil.GenRandomStr()
getSchema := func(colName string, hasPartitionKey bool) *schemapb.CollectionSchema {
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type["fvec_field"] = schemapb.DataType_FloatVector
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
fieldName2Type["int64_field"] = schemapb.DataType_Int64
schema := constructCollectionSchemaByDataType(colName, fieldName2Type, "int64_field", false)
if hasPartitionKey {
partitionKeyField := &schemapb.FieldSchema{
Name: "partition_key_field",
DataType: schemapb.DataType_Int64,
IsPartitionKey: true,
}
fieldName2Type["partition_key_field"] = schemapb.DataType_Int64
schema.Fields = append(schema.Fields, partitionKeyField)
}
return schema
}
getCollectionTask := func(colName string, isIso bool, marshaledSchema []byte) *createCollectionTask {
isoStr := "false"
if isIso {
isoStr = "true"
}
return &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: "",
CollectionName: colName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
Properties: []*commonpb.KeyValuePair{{Key: common.PartitionKeyIsolationKey, Value: isoStr}},
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
}
createIsoCollection := func(colName string, hasPartitionKey bool, isIsolation bool, isIsoNil bool) {
isoStr := "false"
if isIsolation {
isoStr = "true"
}
schema := getSchema(colName, hasPartitionKey)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DropCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: colName,
Schema: marshaledSchema,
ShardsNum: 1,
Properties: []*commonpb.KeyValuePair{{Key: common.PartitionKeyIsolationKey, Value: isoStr}},
}
if isIsoNil {
createColReq.Properties = nil
}
stats, err := qc.CreateCollection(ctx, createColReq)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, stats.ErrorCode)
}
getAlterCollectionTask := func(colName string, isIsolation bool) *alterCollectionTask {
isoStr := "false"
if isIsolation {
isoStr = "true"
}
return &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
Properties: []*commonpb.KeyValuePair{{Key: common.PartitionKeyIsolationKey, Value: isoStr}},
},
mixCoord: qc,
}
}
t.Run("create collection valid", func(t *testing.T) {
paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("true")
defer paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("false")
schema := getSchema(collectionName, true)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := getCollectionTask(collectionName, true, marshaledSchema)
err = createCollectionTask.PreExecute(ctx)
assert.NoError(t, err)
err = createCollectionTask.Execute(ctx)
assert.NoError(t, err)
})
t.Run("create collection without isolation", func(t *testing.T) {
schema := getSchema(collectionName, true)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := getCollectionTask(collectionName, false, marshaledSchema)
err = createCollectionTask.PreExecute(ctx)
assert.NoError(t, err)
err = createCollectionTask.Execute(ctx)
assert.NoError(t, err)
})
t.Run("create collection isolation but no partition key", func(t *testing.T) {
schema := getSchema(collectionName, false)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := getCollectionTask(collectionName, true, marshaledSchema)
assert.ErrorContains(t, createCollectionTask.PreExecute(ctx), "partition key isolation mode is enabled but no partition key field is set")
})
t.Run("create collection with isolation and partition key but MV is not enabled", func(t *testing.T) {
paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("false")
schema := getSchema(collectionName, true)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createCollectionTask := getCollectionTask(collectionName, true, marshaledSchema)
assert.ErrorContains(t, createCollectionTask.PreExecute(ctx), "partition key isolation mode is enabled but current Milvus does not support it")
})
t.Run("alter collection from valid", func(t *testing.T) {
paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("true")
defer paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("false")
colName := collectionName + "AlterValid"
createIsoCollection(colName, true, false, false)
alterTask := getAlterCollectionTask(colName, true)
err := alterTask.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("alter collection without isolation property should succeed", func(t *testing.T) {
colName := collectionName + "AlterNoIso"
createIsoCollection(colName, true, false, true)
alterTask := alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
Properties: nil,
},
mixCoord: qc,
}
err := alterTask.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("alter collection isolation but no partition key", func(t *testing.T) {
colName := collectionName + "AlterNoPartkey"
createIsoCollection(colName, false, false, false)
alterTask := getAlterCollectionTask(colName, true)
assert.ErrorContains(t, alterTask.PreExecute(ctx), "partition key isolation mode is enabled but no partition key field is set")
})
t.Run("alter collection with isolation and partition key but MV is not enabled", func(t *testing.T) {
paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("false")
colName := collectionName + "AlterNoMv"
createIsoCollection(colName, true, false, false)
alterTask := getAlterCollectionTask(colName, true)
assert.ErrorContains(t, alterTask.PreExecute(ctx), "partition key isolation mode is enabled but current Milvus does not support it")
})
t.Run("alter isolation with vector index should fail", func(t *testing.T) {
paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("true")
defer paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("false")
colName := collectionName + "AlterVecIndex"
createIsoCollection(colName, true, true, false)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := getAlterCollectionTask(colName, false)
assert.ErrorContains(t, alterTask.PreExecute(ctx),
"can not alter partition key isolation mode if the collection already has a vector index. Please drop the index first")
})
t.Run("delete isolation property with vector index should fail", func(t *testing.T) {
paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("true")
defer paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("false")
colName := collectionName + "DeleteIsoWithIdx"
createIsoCollection(colName, true, true, false)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
DeleteKeys: []string{common.PartitionKeyIsolationKey},
},
mixCoord: qc,
}
err = alterTask.PreExecute(ctx)
assert.ErrorContains(t, err,
"can not alter partition key isolation mode if the collection already has a vector index. Please drop the index first")
})
t.Run("delete isolation property without vector index should succeed", func(t *testing.T) {
colName := collectionName + "DeleteIsoNoIdx"
createIsoCollection(colName, true, true, false)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
DeleteKeys: []string{common.PartitionKeyIsolationKey},
},
mixCoord: qc,
}
err := alterTask.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("alter unrelated property with isolation enabled should not check index", func(t *testing.T) {
paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("true")
defer paramtable.Get().CommonCfg.EnableMaterializedView.SwapTempValue("false")
colName := collectionName + "AlterOtherPropWithIso"
createIsoCollection(colName, true, true, false)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
Properties: []*commonpb.KeyValuePair{{Key: common.CollectionTTLConfigKey, Value: "3600"}},
},
mixCoord: qc,
}
err = alterTask.PreExecute(ctx)
assert.NoError(t, err)
})
}
func TestAlterCollectionQueryMode(t *testing.T) {
qc := NewMixCoordMock()
ctx := context.Background()
err := InitMetaCache(ctx, qc)
assert.NoError(t, err)
prefix := "TestQueryMode"
getSchema := func(colName string) *schemapb.CollectionSchema {
fieldName2Type := make(map[string]schemapb.DataType)
fieldName2Type["fvec_field"] = schemapb.DataType_FloatVector
fieldName2Type["int64_field"] = schemapb.DataType_Int64
return constructCollectionSchemaByDataType(colName, fieldName2Type, "int64_field", false)
}
createCollection := func(colName string, largeTopKEnabled bool) {
schema := getSchema(colName)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
props := []*commonpb.KeyValuePair{}
if largeTopKEnabled {
props = append(props, &commonpb.KeyValuePair{Key: common.QueryModeKey, Value: common.QueryModeLargeTopK})
}
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_CreateCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: colName,
Schema: marshaledSchema,
ShardsNum: 1,
Properties: props,
}
stats, err := qc.CreateCollection(ctx, createColReq)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, stats.ErrorCode)
}
t.Run("set query_mode to large_topk without vector index", func(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createCollection(colName, false)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
Properties: []*commonpb.KeyValuePair{{Key: common.QueryModeKey, Value: common.QueryModeLargeTopK}},
},
mixCoord: qc,
}
err := alterTask.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("delete query_mode from collection without vector index", func(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createCollection(colName, true)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
DeleteKeys: []string{common.QueryModeKey},
},
mixCoord: qc,
}
err := alterTask.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("set query_mode to large_topk with vector index should fail", func(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createCollection(colName, false)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
Properties: []*commonpb.KeyValuePair{{Key: common.QueryModeKey, Value: common.QueryModeLargeTopK}},
},
mixCoord: qc,
}
err = alterTask.PreExecute(ctx)
assert.ErrorContains(t, err,
"can not alter query_mode if the collection already has a vector index. Please drop the index first")
})
t.Run("delete query_mode with vector index should fail", func(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createCollection(colName, true)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
DeleteKeys: []string{common.QueryModeKey},
},
mixCoord: qc,
}
err = alterTask.PreExecute(ctx)
assert.ErrorContains(t, err,
"can not alter query_mode if the collection already has a vector index. Please drop the index first")
})
t.Run("delete query_mode property without vector index should succeed", func(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createCollection(colName, true)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
DeleteKeys: []string{common.QueryModeKey},
},
mixCoord: qc,
}
err := alterTask.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("alter unrelated property with large_topk query mode should not check index", func(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createCollection(colName, true)
mockVectorIndexForCollection(t, ctx, qc, colName)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
Properties: []*commonpb.KeyValuePair{{Key: common.CollectionTTLConfigKey, Value: "3600"}},
},
mixCoord: qc,
}
err = alterTask.PreExecute(ctx)
assert.NoError(t, err)
})
}
func TestCollectionNamespaceShardingEnabledValidation(t *testing.T) {
qc := NewMixCoordMock()
ctx := context.Background()
err := InitMetaCache(ctx, qc)
assert.NoError(t, err)
prefix := "TestNamespaceShardingEnabled"
getSchemaBytes := func(colName string) []byte {
fieldName2Type := map[string]schemapb.DataType{
"fvec_field": schemapb.DataType_FloatVector,
"int64_field": schemapb.DataType_Int64,
}
schema := constructCollectionSchemaByDataType(colName, fieldName2Type, "int64_field", false)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
return marshaledSchema
}
t.Run("create rejects invalid namespace.sharding.enabled", func(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgID: UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt()),
Timestamp: Timestamp(time.Now().UnixNano()),
},
DbName: "",
CollectionName: colName,
Schema: getSchemaBytes(colName),
ShardsNum: 1,
Properties: []*commonpb.KeyValuePair{{Key: common.NamespaceShardingEnabledKey, Value: "invalid"}},
},
ctx: ctx,
mixCoord: qc,
result: nil,
schema: nil,
}
err := createTask.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid namespace.sharding.enabled")
})
t.Run("alter rejects namespace.sharding.enabled", func(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_CreateCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: colName,
Schema: getSchemaBytes(colName),
ShardsNum: 1,
}
stats, err := qc.CreateCollection(ctx, createColReq)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, stats.ErrorCode)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
Properties: []*commonpb.KeyValuePair{{Key: common.NamespaceShardingEnabledKey, Value: "true"}},
},
mixCoord: qc,
}
err = alterTask.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "cannot alter namespace.sharding.enabled")
})
t.Run("alter rejects deleting namespace.sharding.enabled", func(t *testing.T) {
colName := prefix + funcutil.GenRandomStr()
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_CreateCollection,
MsgID: 101,
Timestamp: 101,
},
DbName: dbName,
CollectionName: colName,
Schema: getSchemaBytes(colName),
ShardsNum: 1,
}
stats, err := qc.CreateCollection(ctx, createColReq)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, stats.ErrorCode)
alterTask := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: colName,
DeleteKeys: []string{common.NamespaceShardingEnabledKey},
},
mixCoord: qc,
}
err = alterTask.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "cannot delete namespace.sharding.enabled")
})
}
func TestAlterCollectionFieldCheckLoaded(t *testing.T) {
qc := NewMixCoordMock()
InitMetaCache(context.Background(), qc)
collectionName := "test_alter_collection_field_check_loaded"
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DropCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
Schema: nil,
ShardsNum: 1,
}
qc.CreateCollection(context.Background(), createColReq)
resp, err := qc.DescribeCollection(context.Background(), &milvuspb.DescribeCollectionRequest{CollectionName: collectionName})
assert.NoError(t, err)
qc.ShowLoadCollectionsFunc = func(ctx context.Context, req *querypb.ShowCollectionsRequest, opts ...grpc.CallOption) (*querypb.ShowCollectionsResponse, error) {
return &querypb.ShowCollectionsResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
CollectionIDs: []int64{resp.CollectionID},
InMemoryPercentages: []int64{100},
}, nil
}
// update property "mmap.enabled" but the collection is loaded
task := &alterCollectionFieldTask{
AlterCollectionFieldRequest: &milvuspb.AlterCollectionFieldRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
Properties: []*commonpb.KeyValuePair{{Key: common.MmapEnabledKey, Value: "true"}},
},
mixCoord: qc,
}
err = task.PreExecute(context.Background())
assert.Equal(t, merr.Code(merr.ErrCollectionLoaded), merr.Code(err))
// delete property "mmap.enabled" but the collection is loaded
task = &alterCollectionFieldTask{
AlterCollectionFieldRequest: &milvuspb.AlterCollectionFieldRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
DeleteKeys: []string{common.MmapEnabledKey},
},
mixCoord: qc,
}
err = task.PreExecute(context.Background())
assert.Equal(t, merr.Code(merr.ErrCollectionLoaded), merr.Code(err))
}
func TestAlterCollectionField(t *testing.T) {
paramtable.Init()
qc := NewMixCoordMock()
InitMetaCache(context.Background(), qc)
collectionName := "test_alter_collection_field"
// Create collection with string and array fields
schema := &schemapb.CollectionSchema{
Name: collectionName,
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "string_field",
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{{Key: "max_length", Value: "100"}},
},
{
FieldID: 101,
Name: "array_field",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
TypeParams: []*commonpb.KeyValuePair{{Key: "max_capacity", Value: "100"}},
},
},
}
schemaBytes, err := proto.Marshal(schema)
assert.NoError(t, err)
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_CreateCollection,
MsgID: 100,
Timestamp: 100,
},
DbName: dbName,
CollectionName: collectionName,
Schema: schemaBytes,
ShardsNum: 1,
}
qc.CreateCollection(context.Background(), createColReq)
// The collection is not loaded
qc.ShowLoadCollectionsFunc = func(ctx context.Context, req *querypb.ShowCollectionsRequest, opts ...grpc.CallOption) (*querypb.ShowCollectionsResponse, error) {
return &querypb.ShowCollectionsResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
CollectionIDs: []int64{},
InMemoryPercentages: []int64{},
}, nil
}
// Test cases
// 1. the collection is not loaded, updating properties is allowed, deleting mmap.enabled/mmap_enabled is allowed
// 2. max_length can only be updated on varchar field or arrar field with varchar element
// 3. max_capacity can only be updated on array field
// 4. invalid number for max_length/max_capacity is not allowed
// 5. not allow to delete max_length/max_capacity
tests := []struct {
name string
fieldName string
properties []*commonpb.KeyValuePair
deleteKeys []string
expectError bool
errCode int32
}{
{
name: "update string field max_length",
fieldName: "string_field",
properties: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "200"},
},
expectError: false,
},
{
name: "update array field max_capacity",
fieldName: "array_field",
properties: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "200"},
},
expectError: false,
},
{
name: "update mmap_enabled",
fieldName: "string_field",
properties: []*commonpb.KeyValuePair{
{Key: common.MmapEnabledKey, Value: "true"},
},
expectError: false,
},
{
name: "invalid property key",
fieldName: "string_field",
properties: []*commonpb.KeyValuePair{
{Key: "invalid_key", Value: "value"},
},
expectError: true,
errCode: merr.Code(merr.ErrParameterInvalid),
},
{
name: "invalid max_length value type",
fieldName: "string_field",
properties: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "not_a_number"},
},
expectError: true,
errCode: merr.Code(merr.ErrParameterInvalid),
},
{
name: "invalid max_capacity value type",
fieldName: "array_field",
properties: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "not_a_number"},
},
expectError: true,
errCode: merr.Code(merr.ErrParameterInvalid),
},
{
name: "max_length exceeds limit",
fieldName: "string_field",
properties: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "65536"},
},
expectError: true,
errCode: merr.Code(merr.ErrParameterInvalid),
},
{
name: "max_capacity exceeds limit",
fieldName: "array_field",
properties: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "5000"},
},
expectError: true,
errCode: merr.Code(merr.ErrParameterInvalid),
},
{
name: "max_capacity invalid range",
fieldName: "array_field",
properties: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "0"},
},
expectError: true,
errCode: merr.Code(merr.ErrParameterInvalid),
},
{
name: "max_capacity invalid type",
fieldName: "string_field",
properties: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "3"},
},
expectError: true,
errCode: merr.Code(merr.ErrParameterInvalid),
},
{
name: "not allow to update max_length on non-varchar field",
fieldName: "array_field",
properties: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "10"},
},
expectError: true,
errCode: merr.Code(merr.ErrParameterInvalid),
},
{
name: "delete mmap.enabled is allowed",
fieldName: "string_field",
deleteKeys: []string{common.MmapEnabledKey},
expectError: false,
},
{
name: "delete mmap_enabled is allowed",
fieldName: "string_field",
deleteKeys: []string{MmapEnabledKey},
expectError: false,
},
{
name: "delete max_length is not allowed",
fieldName: "string_field",
deleteKeys: []string{common.MaxLengthKey},
expectError: true,
errCode: merr.Code(merr.ErrParameterInvalid),
},
{
name: "delete max_capacity is not allowed",
fieldName: "array_field",
deleteKeys: []string{common.MaxCapacityKey},
expectError: true,
errCode: merr.Code(merr.ErrParameterInvalid),
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
task := &alterCollectionFieldTask{
AlterCollectionFieldRequest: &milvuspb.AlterCollectionFieldRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
FieldName: test.fieldName,
Properties: test.properties,
DeleteKeys: test.deleteKeys,
},
mixCoord: qc,
}
err := task.PreExecute(context.Background())
if test.expectError {
assert.Error(t, err)
if test.errCode != 0 {
assert.Equal(t, test.errCode, merr.Code(err))
}
} else {
assert.NoError(t, err)
}
})
}
t.Run("update array field max_capacity with custom limit", func(t *testing.T) {
err := paramtable.Get().Save(paramtable.Get().ProxyCfg.MaxArrayCapacity.Key, "5000")
assert.NoError(t, err)
defer paramtable.Get().Reset(paramtable.Get().ProxyCfg.MaxArrayCapacity.Key)
task := &alterCollectionFieldTask{
AlterCollectionFieldRequest: &milvuspb.AlterCollectionFieldRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
FieldName: "array_field",
Properties: []*commonpb.KeyValuePair{
{Key: common.MaxCapacityKey, Value: "5000"},
},
},
mixCoord: qc,
}
err = task.PreExecute(context.Background())
assert.NoError(t, err)
})
}
// constructCollectionSchemaWithStructArrayField constructs a collection schema specifically for testing StructArrayField
func constructCollectionSchemaWithStructArrayField(collectionName string, structArrayFieldName string, autoID bool) *schemapb.CollectionSchema {
// Primary key field
pkField := &schemapb.FieldSchema{
FieldID: 100,
Name: testInt64Field,
IsPrimaryKey: true,
Description: "primary key field",
DataType: schemapb.DataType_Int64,
AutoID: autoID,
}
// Vector field
vecField := &schemapb.FieldSchema{
FieldID: 101,
Name: testFloatVecField,
IsPrimaryKey: false,
Description: "float vector field",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(testVecDim),
},
},
}
// StructArrayField with various sub-fields for comprehensive testing
structArrayField := &schemapb.StructArrayFieldSchema{
FieldID: 102,
Name: structArrayFieldName,
Description: "struct array field for testing",
Fields: []*schemapb.FieldSchema{
{
FieldID: 1021,
Name: "sub_text_array",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxLengthKey,
Value: strconv.Itoa(testMaxVarCharLength),
},
{
Key: common.MaxCapacityKey,
Value: "50",
},
},
},
{
FieldID: 1022,
Name: "sub_int_array",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int32,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxCapacityKey,
Value: "50",
},
},
},
{
FieldID: 1023,
Name: "sub_float_vector_array",
DataType: schemapb.DataType_ArrayOfVector,
ElementType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(testVecDim),
},
{
Key: common.MaxCapacityKey,
Value: "50",
},
},
},
},
}
return &schemapb.CollectionSchema{
Name: collectionName,
Description: "test collection with struct array field",
AutoID: autoID,
Fields: []*schemapb.FieldSchema{pkField, vecField},
StructArrayFields: []*schemapb.StructArrayFieldSchema{structArrayField},
}
}
// TestCreateCollectionTaskWithStructArrayField tests creating collections with StructArrayField
func TestCreateCollectionTaskWithStructArrayField(t *testing.T) {
mix := NewMixCoordMock()
cache := globalMetaCache
defer func() { globalMetaCache = cache }()
mockCache := NewMockCache(t)
mockCache.EXPECT().GetDatabaseInfo(mock.Anything, mock.Anything).Return(&databaseInfo{}, nil).Maybe()
globalMetaCache = mockCache
ctx := context.Background()
shardsNum := common.DefaultShardsNum
prefix := "TestCreateCollectionTaskWithStructArrayField"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
// Test with StructArrayField
schema := constructCollectionSchemaWithStructArrayField(collectionName, testStructArrayField, false)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
task := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: mix,
result: nil,
schema: nil,
}
t.Run("create collection with struct array field", func(t *testing.T) {
err := task.OnEnqueue()
assert.NoError(t, err)
err = task.PreExecute(ctx)
assert.NoError(t, err)
// Verify schema contains StructArrayFields
assert.NotNil(t, task.schema)
assert.Len(t, task.schema.StructArrayFields, 1)
structArrayField := task.schema.StructArrayFields[0]
assert.Equal(t, testStructArrayField, structArrayField.Name)
assert.Len(t, structArrayField.Fields, 3)
// Verify sub-fields in StructArrayField have been transformed to structName[fieldName] format
subFields := structArrayField.Fields
// After PreExecute, field names should be transformed to structName[fieldName] format
// sub_text_array -> testStructArrayField[sub_text_array]
assert.Equal(t, fmt.Sprintf("%s[sub_text_array]", testStructArrayField), subFields[0].Name)
assert.Equal(t, schemapb.DataType_Array, subFields[0].DataType)
assert.Equal(t, schemapb.DataType_VarChar, subFields[0].ElementType)
// sub_int_array -> testStructArrayField[sub_int_array]
assert.Equal(t, fmt.Sprintf("%s[sub_int_array]", testStructArrayField), subFields[1].Name)
assert.Equal(t, schemapb.DataType_Array, subFields[1].DataType)
assert.Equal(t, schemapb.DataType_Int32, subFields[1].ElementType)
// sub_float_vector_array -> testStructArrayField[sub_float_vector_array]
assert.Equal(t, fmt.Sprintf("%s[sub_float_vector_array]", testStructArrayField), subFields[2].Name)
assert.Equal(t, schemapb.DataType_ArrayOfVector, subFields[2].DataType)
assert.Equal(t, schemapb.DataType_FloatVector, subFields[2].ElementType)
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, task.result.ErrorCode)
err = task.PostExecute(ctx)
assert.NoError(t, err)
})
t.Run("test struct field name transformation with duplicate field names", func(t *testing.T) {
// Create a schema with multiple structs having the same sub-field names
schemaWithDuplicateNames := &schemapb.CollectionSchema{
Name: collectionName + "_duplicate_names",
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: testInt64Field,
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
AutoID: false,
},
{
FieldID: 101,
Name: testFloatVecField,
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: strconv.Itoa(testVecDim),
},
},
},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 102,
Name: "struct1",
Description: "first struct",
Fields: []*schemapb.FieldSchema{
{
FieldID: 1021,
Name: "field_name", // Duplicate name
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int32,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxCapacityKey,
Value: "100",
},
},
},
{
FieldID: 1022,
Name: "common_field", // Another duplicate name
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxLengthKey,
Value: "100",
},
{
Key: common.MaxCapacityKey,
Value: "100",
},
},
},
},
},
{
FieldID: 103,
Name: "struct2",
Description: "second struct",
Fields: []*schemapb.FieldSchema{
{
FieldID: 1031,
Name: "field_name", // Same name as struct1's field
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Float,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxCapacityKey,
Value: "100",
},
},
},
{
FieldID: 1032,
Name: "common_field", // Same name as struct1's field
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Bool,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.MaxCapacityKey,
Value: "100",
},
},
},
},
},
},
}
marshaledDuplicateSchema, err := proto.Marshal(schemaWithDuplicateNames)
assert.NoError(t, err)
duplicateTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName + "_duplicate_names",
Schema: marshaledDuplicateSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: mix,
result: nil,
schema: nil,
}
err = duplicateTask.OnEnqueue()
assert.NoError(t, err)
// PreExecute should succeed and transform the field names
err = duplicateTask.PreExecute(ctx)
assert.NoError(t, err)
// Verify that field names have been transformed to avoid conflicts
assert.NotNil(t, duplicateTask.schema)
assert.Len(t, duplicateTask.schema.StructArrayFields, 2)
// Check struct1's fields
struct1 := duplicateTask.schema.StructArrayFields[0]
assert.Equal(t, "struct1", struct1.Name)
assert.Equal(t, "struct1[field_name]", struct1.Fields[0].Name)
assert.Equal(t, "struct1[common_field]", struct1.Fields[1].Name)
// Check struct2's fields
struct2 := duplicateTask.schema.StructArrayFields[1]
assert.Equal(t, "struct2", struct2.Name)
assert.Equal(t, "struct2[field_name]", struct2.Fields[0].Name)
assert.Equal(t, "struct2[common_field]", struct2.Fields[1].Name)
// Verify that despite having the same original names, the transformed names are unique
allFieldNames := make(map[string]bool)
for _, structField := range duplicateTask.schema.StructArrayFields {
for _, field := range structField.Fields {
// Each transformed name should be unique
assert.False(t, allFieldNames[field.Name], "Duplicate field name found: %s", field.Name)
allFieldNames[field.Name] = true
}
}
})
t.Run("validate struct array field constraints", func(t *testing.T) {
// Test invalid sub-field in StructArrayField
invalidSchema := constructCollectionSchemaWithStructArrayField(collectionName+"_invalid", testStructArrayField, false)
// Add an invalid sub-field (nested array)
invalidSubField := &schemapb.FieldSchema{
FieldID: 1024,
Name: "invalid_nested_array",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Array, // This should be invalid - nested arrays
}
invalidSchema.StructArrayFields[0].Fields = append(invalidSchema.StructArrayFields[0].Fields, invalidSubField)
invalidMarshaledSchema, err := proto.Marshal(invalidSchema)
assert.NoError(t, err)
invalidTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
Base: nil,
DbName: dbName,
CollectionName: collectionName + "_invalid",
Schema: invalidMarshaledSchema,
ShardsNum: shardsNum,
},
ctx: ctx,
mixCoord: mix,
result: nil,
schema: nil,
}
err = invalidTask.OnEnqueue()
assert.NoError(t, err)
err = invalidTask.PreExecute(ctx)
assert.Error(t, err) // Should fail due to nested array validation
})
}
// TestDescribeCollectionTaskWithStructArrayField tests describing collections with StructArrayField
func TestDescribeCollectionTaskWithStructArrayField(t *testing.T) {
mix := NewMixCoordMock()
ctx := context.Background()
prefix := "TestDescribeCollectionTaskWithStructArrayField"
collectionName := prefix + funcutil.GenRandomStr()
// Create collection with StructArrayField first
schema := constructCollectionSchemaWithStructArrayField(collectionName, testStructArrayField, true)
marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err)
createTask := &createCollectionTask{
Condition: NewTaskCondition(ctx),
CreateCollectionRequest: &milvuspb.CreateCollectionRequest{
CollectionName: collectionName,
Schema: marshaledSchema,
},
ctx: ctx,
mixCoord: mix,
}
err = createTask.OnEnqueue()
assert.NoError(t, err)
err = createTask.PreExecute(ctx)
assert.NoError(t, err)
err = createTask.Execute(ctx)
assert.NoError(t, err)
// Now test describe collection
describeTask := &describeCollectionTask{
Condition: NewTaskCondition(ctx),
DescribeCollectionRequest: &milvuspb.DescribeCollectionRequest{
CollectionName: collectionName,
},
ctx: ctx,
mixCoord: mix,
}
t.Run("describe collection with struct array field", func(t *testing.T) {
err := describeTask.PreExecute(ctx)
assert.NoError(t, err)
err = describeTask.Execute(ctx)
assert.NoError(t, err)
assert.NotNil(t, describeTask.result)
assert.Equal(t, commonpb.ErrorCode_Success, describeTask.result.Status.ErrorCode)
// Verify StructArrayFields are returned in describe response
resultSchema := describeTask.result.Schema
assert.NotNil(t, resultSchema)
assert.Len(t, resultSchema.StructArrayFields, 1)
structArrayField := resultSchema.StructArrayFields[0]
assert.Equal(t, testStructArrayField, structArrayField.Name)
assert.Len(t, structArrayField.Fields, 3)
err = describeTask.PostExecute(ctx)
assert.NoError(t, err)
})
}
func TestAlterCollection_AllowInsertAutoID_AutoIDFalse(t *testing.T) {
qc := NewMixCoordMock()
InitMetaCache(context.Background(), qc)
ctx := context.Background()
collectionName := "test_alter_allow_insert_autoid_autoid_false"
schema := constructCollectionSchemaWithStructArrayField(collectionName, testStructArrayField, false)
schemaBytes, err := proto.Marshal(schema)
assert.NoError(t, err)
createColReq := &milvuspb.CreateCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_CreateCollection,
MsgID: 200,
Timestamp: 200,
},
DbName: dbName,
CollectionName: collectionName,
Schema: schemaBytes,
ShardsNum: 1,
}
qc.CreateCollection(ctx, createColReq)
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{},
CollectionName: collectionName,
Properties: []*commonpb.KeyValuePair{
{Key: common.AllowInsertAutoIDKey, Value: "true"},
},
},
mixCoord: qc,
}
err = task.PreExecute(ctx)
assert.Error(t, err)
assert.Equal(t, merr.Code(merr.ErrParameterInvalid), merr.Code(err))
}
func TestHighlightTask(t *testing.T) {
paramtable.Init()
ctx := context.Background()
t.Run("basic methods", func(t *testing.T) {
task := &HighlightTask{
ctx: ctx,
GetHighlightRequest: &querypb.GetHighlightRequest{
Base: commonpbutil.NewMsgBase(),
},
Condition: NewTaskCondition(ctx),
}
t.Run("traceCtx", func(t *testing.T) {
traceCtx := task.TraceCtx()
assert.NotNil(t, traceCtx)
assert.Equal(t, ctx, traceCtx)
})
t.Run("id", func(t *testing.T) {
id := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
task.SetID(id)
assert.Equal(t, id, task.ID())
})
t.Run("name", func(t *testing.T) {
assert.Equal(t, HighlightTaskName, task.Name())
})
t.Run("type", func(t *testing.T) {
assert.Equal(t, commonpb.MsgType_Undefined, task.Type())
})
t.Run("ts", func(t *testing.T) {
ts := Timestamp(time.Now().UnixNano())
task.SetTs(ts)
assert.Equal(t, ts, task.BeginTs())
assert.Equal(t, ts, task.EndTs())
})
})
t.Run("OnEnqueue", func(t *testing.T) {
task := &HighlightTask{
ctx: ctx,
GetHighlightRequest: &querypb.GetHighlightRequest{},
Condition: NewTaskCondition(ctx),
}
err := task.OnEnqueue()
assert.NoError(t, err)
assert.NotNil(t, task.Base)
assert.Equal(t, commonpb.MsgType_Undefined, task.Base.MsgType)
assert.Equal(t, paramtable.GetNodeID(), task.Base.SourceID)
})
t.Run("PreExecute", func(t *testing.T) {
task := &HighlightTask{
ctx: ctx,
GetHighlightRequest: &querypb.GetHighlightRequest{
Base: commonpbutil.NewMsgBase(),
},
Condition: NewTaskCondition(ctx),
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("getHighlightOnShardleader success", func(t *testing.T) {
task := &HighlightTask{
ctx: ctx,
GetHighlightRequest: &querypb.GetHighlightRequest{
Base: commonpbutil.NewMsgBase(),
Topks: []int64{10},
Tasks: []*querypb.HighlightTask{
{},
},
},
Condition: NewTaskCondition(ctx),
}
mockQN := mocks.NewMockQueryNodeClient(t)
expectedResp := &querypb.GetHighlightResponse{
Status: merr.Success(),
Results: []*querypb.HighlightResult{
{
Fragments: []*querypb.HighlightFragment{
{
StartOffset: 0,
EndOffset: 10,
Offsets: []int64{0, 5, 5, 10},
},
},
},
},
}
mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(expectedResp, nil)
err := task.getHighlightOnShardleader(ctx, 1, mockQN, "test_channel")
assert.NoError(t, err)
assert.NotNil(t, task.result)
assert.Equal(t, expectedResp, task.result)
assert.Equal(t, "test_channel", task.Channel)
})
t.Run("getHighlightOnShardleader rpc error", func(t *testing.T) {
task := &HighlightTask{
ctx: ctx,
GetHighlightRequest: &querypb.GetHighlightRequest{
Base: commonpbutil.NewMsgBase(),
},
Condition: NewTaskCondition(ctx),
}
mockQN := mocks.NewMockQueryNodeClient(t)
mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(nil, errors.New("rpc error"))
err := task.getHighlightOnShardleader(ctx, 1, mockQN, "test_channel")
assert.Error(t, err)
assert.Contains(t, err.Error(), "rpc error")
})
t.Run("getHighlightOnShardleader status error", func(t *testing.T) {
task := &HighlightTask{
ctx: ctx,
GetHighlightRequest: &querypb.GetHighlightRequest{
Base: commonpbutil.NewMsgBase(),
},
Condition: NewTaskCondition(ctx),
}
mockQN := mocks.NewMockQueryNodeClient(t)
expectedResp := &querypb.GetHighlightResponse{
Status: merr.Status(errors.New("status error")),
}
mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(expectedResp, nil)
err := task.getHighlightOnShardleader(ctx, 1, mockQN, "test_channel")
assert.Error(t, err)
})
t.Run("Execute success", func(t *testing.T) {
task := &HighlightTask{
ctx: ctx,
collectionName: "test_collection",
collectionID: 100,
dbName: "default",
GetHighlightRequest: &querypb.GetHighlightRequest{
Base: commonpbutil.NewMsgBase(),
Topks: []int64{10, 20},
Tasks: []*querypb.HighlightTask{
{},
{},
},
},
Condition: NewTaskCondition(ctx),
}
mockLB := shardclient.NewMockLBPolicy(t)
mockQN := mocks.NewMockQueryNodeClient(t)
expectedResp := &querypb.GetHighlightResponse{
Status: merr.Success(),
Results: []*querypb.HighlightResult{
{
Fragments: []*querypb.HighlightFragment{
{
StartOffset: 0,
EndOffset: 10,
Offsets: []int64{0, 5, 5, 10},
},
},
},
},
}
mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(expectedResp, nil)
mockLB.EXPECT().ExecuteOneChannel(mock.Anything, mock.Anything).Run(
func(ctx context.Context, workload shardclient.CollectionWorkLoad) {
assert.Equal(t, "default", workload.Db)
assert.Equal(t, "test_collection", workload.CollectionName)
assert.Equal(t, int64(100), workload.CollectionID)
assert.Equal(t, int64(4), workload.Nq) // len(topks) * len(tasks) = 2 * 2 = 4
err := workload.Exec(ctx, 1, mockQN, "test_channel")
assert.NoError(t, err)
},
).Return(nil)
task.lb = mockLB
err := task.Execute(ctx)
assert.NoError(t, err)
assert.NotNil(t, task.result)
})
t.Run("Execute lb error", func(t *testing.T) {
task := &HighlightTask{
ctx: ctx,
collectionName: "test_collection",
collectionID: 100,
dbName: "default",
GetHighlightRequest: &querypb.GetHighlightRequest{
Base: commonpbutil.NewMsgBase(),
},
Condition: NewTaskCondition(ctx),
}
mockLB := shardclient.NewMockLBPolicy(t)
mockLB.EXPECT().ExecuteOneChannel(mock.Anything, mock.Anything).Return(errors.New("lb error"))
task.lb = mockLB
err := task.Execute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "lb error")
})
t.Run("PostExecute", func(t *testing.T) {
task := &HighlightTask{
ctx: ctx,
GetHighlightRequest: &querypb.GetHighlightRequest{
Base: commonpbutil.NewMsgBase(),
},
Condition: NewTaskCondition(ctx),
}
err := task.PostExecute(ctx)
assert.NoError(t, err)
})
}
func TestValidateCollectionTTL(t *testing.T) {
cases := []struct {
name string
value string
hasTTL bool
expectErr bool
}{
{"at_min_boundary", "-1", true, false}, // val <= -1
{"out_of_lower_boundary", "-2", true, true}, // val < -1
{"exceed_max", "3155760001", true, true}, // val > MaxTTLSeconds
{"at_max_boundary", "3155760000", true, false}, // val == MaxTTLSeconds
{"valid_value", "3600", true, false},
{"min_valid", "1", true, false},
{"invalid_format", "abc", true, true},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
hasTTL, err := validateCollectionTTL([]*commonpb.KeyValuePair{{Key: common.CollectionTTLConfigKey, Value: c.value}})
assert.Equal(t, c.hasTTL, hasTTL)
if c.expectErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestHasWarmupProp(t *testing.T) {
t.Run("has generic warmup key", func(t *testing.T) {
props := []*commonpb.KeyValuePair{
{Key: common.WarmupKey, Value: common.WarmupSync},
}
assert.True(t, hasWarmupProp(props...))
})
t.Run("has scalar field warmup key", func(t *testing.T) {
props := []*commonpb.KeyValuePair{
{Key: common.WarmupScalarFieldKey, Value: common.WarmupDisable},
}
assert.True(t, hasWarmupProp(props...))
})
t.Run("has scalar index warmup key", func(t *testing.T) {
props := []*commonpb.KeyValuePair{
{Key: common.WarmupScalarIndexKey, Value: common.WarmupSync},
}
assert.True(t, hasWarmupProp(props...))
})
t.Run("has vector field warmup key", func(t *testing.T) {
props := []*commonpb.KeyValuePair{
{Key: common.WarmupVectorFieldKey, Value: common.WarmupDisable},
}
assert.True(t, hasWarmupProp(props...))
})
t.Run("has vector index warmup key", func(t *testing.T) {
props := []*commonpb.KeyValuePair{
{Key: common.WarmupVectorIndexKey, Value: common.WarmupSync},
}
assert.True(t, hasWarmupProp(props...))
})
t.Run("no warmup key", func(t *testing.T) {
props := []*commonpb.KeyValuePair{
{Key: "other_key", Value: "other_value"},
}
assert.False(t, hasWarmupProp(props...))
})
t.Run("empty props", func(t *testing.T) {
assert.False(t, hasWarmupProp())
})
}
func TestHasPropInDeletekeys(t *testing.T) {
t.Run("has mmap key", func(t *testing.T) {
keys := []string{common.MmapEnabledKey}
assert.Equal(t, common.MmapEnabledKey, hasPropInDeletekeys(keys))
})
t.Run("has generic warmup key", func(t *testing.T) {
keys := []string{common.WarmupKey}
assert.Equal(t, common.WarmupKey, hasPropInDeletekeys(keys))
})
t.Run("has granular warmup keys", func(t *testing.T) {
assert.Equal(t, common.WarmupScalarFieldKey, hasPropInDeletekeys([]string{common.WarmupScalarFieldKey}))
assert.Equal(t, common.WarmupScalarIndexKey, hasPropInDeletekeys([]string{common.WarmupScalarIndexKey}))
assert.Equal(t, common.WarmupVectorFieldKey, hasPropInDeletekeys([]string{common.WarmupVectorFieldKey}))
assert.Equal(t, common.WarmupVectorIndexKey, hasPropInDeletekeys([]string{common.WarmupVectorIndexKey}))
})
t.Run("no special key", func(t *testing.T) {
keys := []string{"other_key"}
assert.Equal(t, "", hasPropInDeletekeys(keys))
})
t.Run("empty keys", func(t *testing.T) {
assert.Equal(t, "", hasPropInDeletekeys([]string{}))
})
}
func TestValidateAddFieldRequest(t *testing.T) {
// Build a base schema with one int64 PK and one float vector field.
baseSchema := func() *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Name: "test_coll",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{
FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}},
},
},
}
}
t.Run("happy path", func(t *testing.T) {
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "new_field",
DataType: schemapb.DataType_Int64,
Nullable: true,
}
err := validateAddFieldRequest(schema, newField)
assert.NoError(t, err)
})
t.Run("text field requires storage v3", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "false")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
})
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "text_field",
DataType: schemapb.DataType_Text,
Nullable: true,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "TEXT field requires StorageV3")
})
t.Run("text field allowed with storage v3", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
})
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "text_field",
DataType: schemapb.DataType_Text,
Nullable: true,
}
err := validateAddFieldRequest(schema, newField)
assert.NoError(t, err)
})
t.Run("text add field does not support default value", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
})
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "text_field",
DataType: schemapb.DataType_Text,
Nullable: true,
DefaultValue: &schemapb.ValueField{
Data: &schemapb.ValueField_StringData{StringData: "default text"},
},
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "default value is not supported when adding TEXT field")
})
t.Run("field count exceeds MaxFieldNum", func(t *testing.T) {
schema := baseSchema()
// Build exactly MaxFieldNum fields so that the check triggers.
maxFieldNum := Params.ProxyCfg.MaxFieldNum.GetAsInt()
for i := len(schema.Fields); i < maxFieldNum; i++ {
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
FieldID: int64(200 + i),
Name: fmt.Sprintf("field_%d", i),
DataType: schemapb.DataType_Int64,
})
}
newField := &schemapb.FieldSchema{
Name: "overflow_field",
DataType: schemapb.DataType_Int64,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("duplicate field name", func(t *testing.T) {
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "vec", // already exists
DataType: schemapb.DataType_Int64,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("invalid DataType_None", func(t *testing.T) {
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "invalid_type_field",
DataType: schemapb.DataType_None,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("external field mapping", func(t *testing.T) {
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "external_field",
DataType: schemapb.DataType_Int64,
ExternalField: "remote_field",
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "does not support external field mapping")
})
t.Run("external collection requires external field mapping", func(t *testing.T) {
schema := baseSchema()
for _, field := range schema.GetFields() {
field.ExternalField = field.GetName()
}
newField := &schemapb.FieldSchema{
Name: "score",
DataType: schemapb.DataType_Double,
Nullable: true,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "requires external_field mapping")
})
t.Run("external collection allows external field mapping", func(t *testing.T) {
schema := baseSchema()
for _, field := range schema.GetFields() {
field.ExternalField = field.GetName()
}
newField := &schemapb.FieldSchema{
Name: "score",
DataType: schemapb.DataType_Double,
Nullable: true,
ExternalField: "score",
}
err := validateAddFieldRequest(schema, newField)
assert.NoError(t, err)
})
t.Run("external collection rejects unsupported field type", func(t *testing.T) {
schema := baseSchema()
for _, field := range schema.GetFields() {
field.ExternalField = field.GetName()
}
newField := &schemapb.FieldSchema{
Name: "sparse",
DataType: schemapb.DataType_SparseFloatVector,
ExternalField: "sparse",
}
err := validateAddFieldRequest(schema, newField)
if assert.Error(t, err) {
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "does not support field type")
}
})
t.Run("external collection rejects duplicate external field mapping", func(t *testing.T) {
schema := baseSchema()
for _, field := range schema.GetFields() {
field.ExternalField = field.GetName()
}
newField := &schemapb.FieldSchema{
Name: "score",
DataType: schemapb.DataType_Double,
Nullable: true,
ExternalField: "vec",
}
err := validateAddFieldRequest(schema, newField)
if assert.Error(t, err) {
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.Contains(t, err.Error(), "external_field \"vec\" is mapped by multiple fields")
}
})
t.Run("system field name RowID", func(t *testing.T) {
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "RowID",
DataType: schemapb.DataType_Int64,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("is_primary_key = true", func(t *testing.T) {
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "new_pk",
DataType: schemapb.DataType_Int64,
IsPrimaryKey: true,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("auto_id = true", func(t *testing.T) {
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "auto_field",
DataType: schemapb.DataType_Int64,
AutoID: true,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("is_partition_key = true", func(t *testing.T) {
schema := baseSchema()
newField := &schemapb.FieldSchema{
Name: "part_key",
DataType: schemapb.DataType_Int64,
IsPartitionKey: true,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("clustering key when schema already has one", func(t *testing.T) {
schema := baseSchema()
// Mark an existing field as clustering key.
schema.Fields[0].IsClusteringKey = true
newField := &schemapb.FieldSchema{
Name: "new_cluster",
DataType: schemapb.DataType_Int64,
IsClusteringKey: true,
}
err := validateAddFieldRequest(schema, newField)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("vector type with max vector fields exceeded", func(t *testing.T) {
schema := baseSchema()
// Fill up vector fields to MaxVectorFieldNum.
maxVecFieldNum := Params.ProxyCfg.MaxVectorFieldNum.GetAsInt()
// schema already has one FloatVector; add more to reach the limit.
for i := 1; i < maxVecFieldNum; i++ {
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
FieldID: int64(300 + i),
Name: fmt.Sprintf("vec_%d", i),
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
},
})
}
newVecField := &schemapb.FieldSchema{
Name: "extra_vec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
},
}
err := validateAddFieldRequest(schema, newVecField)
assert.Error(t, err)
})
}
func TestAlterCollectionSchemaTask(t *testing.T) {
// Add-function cases require StorageV3 + schema-bump/storage-version compaction enabled
// (validateAddFunctionRequiresStorageV3). storageVersion.enabled defaults true;
// bumpSchemaVersion.enabled defaults false, so it must be set explicitly.
paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, "true")
paramtable.Get().Save(paramtable.Get().DataCoordCfg.BumpSchemaVersionCompactionEnabled.Key, "true")
defer paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
defer paramtable.Get().Reset(paramtable.Get().DataCoordCfg.BumpSchemaVersionCompactionEnabled.Key)
rc := NewMixCoordMock()
ctx := context.Background()
prefix := "TestAlterCollectionSchemaTask"
dbName := ""
collectionName := prefix + funcutil.GenRandomStr()
collectionID := int64(1)
rc.collName2ID[collectionName] = collectionID
rc.collID2Meta[collectionID] = collectionMeta{
name: collectionName,
id: collectionID,
schema: &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int64, AutoID: true, Name: "ID", IsPrimaryKey: true},
{
FieldID: 101, DataType: schemapb.DataType_VarChar, Name: "text",
TypeParams: []*commonpb.KeyValuePair{
{Key: "max_length", Value: "65535"},
},
},
},
},
}
// oldSchema matches the collection: has a varchar field "text" used as BM25 input.
oldSchema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, DataType: schemapb.DataType_Int64, AutoID: true, Name: "ID", IsPrimaryKey: true},
{
FieldID: 101, DataType: schemapb.DataType_VarChar, Name: "text",
TypeParams: []*commonpb.KeyValuePair{
{Key: "max_length", Value: "65535"},
{Key: "enable_analyzer", Value: "true"},
},
},
},
}
// sparseOutputField is the output field for the BM25 function.
sparseOutputField := &schemapb.FieldSchema{
Name: "sparse_bm25",
DataType: schemapb.DataType_SparseFloatVector,
}
// BM25 function schema: input "text", output "sparse_bm25".
functionSchema := &schemapb.FunctionSchema{
Name: "bm25_func",
Type: schemapb.FunctionType_BM25,
InputFieldNames: []string{"text"},
OutputFieldNames: []string{"sparse_bm25"},
}
functionOnlyOldSchema := proto.Clone(oldSchema).(*schemapb.CollectionSchema)
functionOnlyOldSchema.Fields = append(functionOnlyOldSchema.Fields, &schemapb.FieldSchema{
FieldID: 102,
Name: "existing_minhash_output",
DataType: schemapb.DataType_BinaryVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "4096"},
},
})
minHashFunctionOnlySchema := &schemapb.FunctionSchema{
Name: "minhash_existing_func",
Type: schemapb.FunctionType_MinHash,
InputFieldNames: []string{"text"},
OutputFieldNames: []string{"existing_minhash_output"},
Params: []*commonpb.KeyValuePair{
{Key: "num_hashes", Value: "128"},
{Key: "shingle_size", Value: "3"},
{Key: "hash_function", Value: "xxhash64"},
{Key: "seed", Value: "42"},
},
}
buildValidRequest := func() *milvuspb.AlterCollectionSchemaRequest {
return &milvuspb.AlterCollectionSchemaRequest{
DbName: dbName,
CollectionName: collectionName,
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_AddRequest{
AddRequest: &milvuspb.AlterCollectionSchemaRequest_AddRequest{
FieldInfos: []*milvuspb.AlterCollectionSchemaRequest_FieldInfo{
{
FieldSchema: proto.Clone(sparseOutputField).(*schemapb.FieldSchema),
ExtraParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "SPARSE_INVERTED_INDEX"},
{Key: common.MetricTypeKey, Value: "BM25"},
},
},
},
FuncSchema: []*schemapb.FunctionSchema{proto.Clone(functionSchema).(*schemapb.FunctionSchema)},
},
},
},
}
}
buildAddFieldRequest := func(fieldName string, nullable bool, doBackfill bool) *milvuspb.AlterCollectionSchemaRequest {
return &milvuspb.AlterCollectionSchemaRequest{
DbName: dbName,
CollectionName: collectionName,
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_AddRequest{
AddRequest: &milvuspb.AlterCollectionSchemaRequest_AddRequest{
FieldInfos: []*milvuspb.AlterCollectionSchemaRequest_FieldInfo{
{FieldSchema: &schemapb.FieldSchema{
Name: fieldName,
DataType: schemapb.DataType_VarChar,
Nullable: nullable,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "128"},
},
}},
},
DoPhysicalBackfill: doBackfill,
},
},
},
}
}
buildAddFunctionRequest := func(functionSchema *schemapb.FunctionSchema) *milvuspb.AlterCollectionSchemaRequest {
return &milvuspb.AlterCollectionSchemaRequest{
DbName: dbName,
CollectionName: collectionName,
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_AddRequest{
AddRequest: &milvuspb.AlterCollectionSchemaRequest_AddRequest{
FuncSchema: []*schemapb.FunctionSchema{functionSchema},
},
},
},
}
}
buildTask := func(req *milvuspb.AlterCollectionSchemaRequest, schema *schemapb.CollectionSchema) *alterCollectionSchemaTask {
return &alterCollectionSchemaTask{
Condition: NewTaskCondition(ctx),
AlterCollectionSchemaRequest: req,
AlterCollectionSchemaResponse: nil,
ctx: ctx,
mixCoord: rc,
oldSchema: schema,
}
}
t.Run("OnEnqueue", func(t *testing.T) {
task := buildTask(buildValidRequest(), oldSchema)
err := task.OnEnqueue()
assert.NoError(t, err)
assert.Equal(t, commonpb.MsgType_AlterCollectionSchema, task.Type())
})
t.Run("ctx", func(t *testing.T) {
task := buildTask(buildValidRequest(), oldSchema)
assert.NotNil(t, task.TraceCtx())
})
t.Run("id", func(t *testing.T) {
task := buildTask(buildValidRequest(), oldSchema)
_ = task.OnEnqueue()
id := UniqueID(uniquegenerator.GetUniqueIntGeneratorIns().GetInt())
task.SetID(id)
assert.Equal(t, id, task.ID())
})
t.Run("name", func(t *testing.T) {
task := buildTask(buildValidRequest(), oldSchema)
assert.Equal(t, AlterCollectionSchemaTaskName, task.Name())
})
t.Run("ts", func(t *testing.T) {
task := buildTask(buildValidRequest(), oldSchema)
_ = task.OnEnqueue()
ts := Timestamp(time.Now().UnixNano())
task.SetTs(ts)
assert.Equal(t, ts, task.BeginTs())
assert.Equal(t, ts, task.EndTs())
})
t.Run("PostExecute", func(t *testing.T) {
task := buildTask(buildValidRequest(), oldSchema)
assert.NoError(t, task.PostExecute(ctx))
})
t.Run("PreExecute nil oldSchema", func(t *testing.T) {
task := buildTask(buildValidRequest(), nil)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("PreExecute nil action", func(t *testing.T) {
req := &milvuspb.AlterCollectionSchemaRequest{
DbName: dbName,
CollectionName: collectionName,
Action: nil,
}
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("PreExecute nil addRequest", func(t *testing.T) {
// Action is set but Op is not AddRequest (use nil oneof).
req := &milvuspb.AlterCollectionSchemaRequest{
DbName: dbName,
CollectionName: collectionName,
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
// Op is nil — GetAddRequest() returns nil.
},
}
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("PreExecute add field without funcSchema", func(t *testing.T) {
task := buildTask(buildAddFieldRequest("extra_text", true, false), oldSchema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("PreExecute add field without funcSchema ignores physical backfill flag", func(t *testing.T) {
task := buildTask(buildAddFieldRequest("extra_text_backfill", true, true), oldSchema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("PreExecute add field without funcSchema requires nullable field", func(t *testing.T) {
task := buildTask(buildAddFieldRequest("extra_text_not_nullable", false, false), oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.ErrorContains(t, err, "added field must be nullable")
})
t.Run("PreExecute rejects duplicated add field name", func(t *testing.T) {
task := buildTask(buildAddFieldRequest("text", true, false), oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.ErrorContains(t, err, "duplicated field name")
})
t.Run("PreExecute rejects invalid add field name", func(t *testing.T) {
task := buildTask(buildAddFieldRequest("invalid field name", true, false), oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorContains(t, err, "Invalid field name")
})
t.Run("PreExecute multiple funcSchemas", func(t *testing.T) {
req := &milvuspb.AlterCollectionSchemaRequest{
DbName: dbName,
CollectionName: collectionName,
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_AddRequest{
AddRequest: &milvuspb.AlterCollectionSchemaRequest_AddRequest{
FieldInfos: []*milvuspb.AlterCollectionSchemaRequest_FieldInfo{
{FieldSchema: sparseOutputField},
},
FuncSchema: []*schemapb.FunctionSchema{functionSchema, functionSchema}, // two functions
},
},
},
}
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("PreExecute nil funcSchema", func(t *testing.T) {
task := buildTask(buildAddFunctionRequest(nil), oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.ErrorContains(t, err, "function schema is nil")
})
t.Run("PreExecute add function without fieldInfos", func(t *testing.T) {
task := buildTask(buildAddFunctionRequest(minHashFunctionOnlySchema), functionOnlyOldSchema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("PreExecute rejects function-only BM25 request", func(t *testing.T) {
task := buildTask(buildAddFunctionRequest(functionSchema), oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.ErrorContains(t, err, "add_function_field")
})
t.Run("PreExecute rejects empty add request", func(t *testing.T) {
req := &milvuspb.AlterCollectionSchemaRequest{
DbName: dbName,
CollectionName: collectionName,
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_AddRequest{
AddRequest: &milvuspb.AlterCollectionSchemaRequest_AddRequest{},
},
},
}
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("PreExecute multiple fieldInfos", func(t *testing.T) {
req := &milvuspb.AlterCollectionSchemaRequest{
DbName: dbName,
CollectionName: collectionName,
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_AddRequest{
AddRequest: &milvuspb.AlterCollectionSchemaRequest_AddRequest{
FieldInfos: []*milvuspb.AlterCollectionSchemaRequest_FieldInfo{
{FieldSchema: sparseOutputField},
{FieldSchema: &schemapb.FieldSchema{Name: "extra", DataType: schemapb.DataType_Int64}},
},
FuncSchema: []*schemapb.FunctionSchema{functionSchema},
},
},
},
}
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("PreExecute nil fieldSchema in fieldInfo", func(t *testing.T) {
req := &milvuspb.AlterCollectionSchemaRequest{
DbName: dbName,
CollectionName: collectionName,
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_AddRequest{
AddRequest: &milvuspb.AlterCollectionSchemaRequest_AddRequest{
FieldInfos: []*milvuspb.AlterCollectionSchemaRequest_FieldInfo{
{FieldSchema: nil}, // nil inner schema
},
FuncSchema: []*schemapb.FunctionSchema{functionSchema},
},
},
},
}
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("PreExecute output field points to existing field", func(t *testing.T) {
// OutputFieldNames must only reference newly-added fields, not pre-existing ones.
// "text" is already in oldSchema; pointing function output to it must be rejected.
req := &milvuspb.AlterCollectionSchemaRequest{
DbName: dbName,
CollectionName: collectionName,
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_AddRequest{
AddRequest: &milvuspb.AlterCollectionSchemaRequest_AddRequest{
FieldInfos: []*milvuspb.AlterCollectionSchemaRequest_FieldInfo{
{FieldSchema: sparseOutputField},
},
FuncSchema: []*schemapb.FunctionSchema{
{
Name: "bm25_func",
Type: schemapb.FunctionType_BM25,
InputFieldNames: []string{"text"},
OutputFieldNames: []string{"text"}, // "text" exists in oldSchema — must be rejected
},
},
},
},
},
}
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("PreExecute supports MinHash function with new output field", func(t *testing.T) {
req := buildValidRequest()
addRequest := req.GetAction().GetAddRequest()
addRequest.FieldInfos[0].FieldSchema = &schemapb.FieldSchema{
Name: "binary_minhash",
DataType: schemapb.DataType_BinaryVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "4096"},
},
}
addRequest.FieldInfos[0].ExtraParams = []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "MINHASH_LSH"},
{Key: common.MetricTypeKey, Value: "MHJACCARD"},
}
functionSchema := proto.Clone(addRequest.GetFuncSchema()[0]).(*schemapb.FunctionSchema)
functionSchema.Name = "minhash_func"
functionSchema.Type = schemapb.FunctionType_MinHash
functionSchema.OutputFieldNames = []string{"binary_minhash"}
functionSchema.Params = []*commonpb.KeyValuePair{
{Key: "num_hashes", Value: "128"},
{Key: "shingle_size", Value: "3"},
{Key: "hash_function", Value: "xxhash64"},
{Key: "seed", Value: "42"},
}
addRequest.FuncSchema = []*schemapb.FunctionSchema{functionSchema}
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("PreExecute rejects unsupported function", func(t *testing.T) {
req := buildValidRequest()
addRequest := req.GetAction().GetAddRequest()
functionSchema := proto.Clone(addRequest.GetFuncSchema()[0]).(*schemapb.FunctionSchema)
functionSchema.Type = schemapb.FunctionType_TextEmbedding
addRequest.FuncSchema = []*schemapb.FunctionSchema{functionSchema}
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.ErrorContains(t, err, "only BM25 and MinHash functions are supported")
})
t.Run("PreExecute ignores legacy DoPhysicalBackfill", func(t *testing.T) {
schema := proto.Clone(oldSchema).(*schemapb.CollectionSchema)
schema.DoPhysicalBackfill = true
task := buildTask(buildValidRequest(), schema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("PreExecute rejects multiple function outputs", func(t *testing.T) {
req := buildValidRequest()
req.GetAction().GetAddRequest().GetFuncSchema()[0].OutputFieldNames = []string{"sparse_bm25", "sparse_bm25_extra"}
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.ErrorContains(t, err, "exactly one output field")
})
t.Run("PreExecute rejects function-only request with missing output field", func(t *testing.T) {
req := buildAddFunctionRequest(&schemapb.FunctionSchema{
Name: "missing_output_func",
Type: schemapb.FunctionType_MinHash,
InputFieldNames: []string{"text"},
OutputFieldNames: []string{"missing_output"},
Params: []*commonpb.KeyValuePair{
{Key: "num_hashes", Value: "128"},
{Key: "shingle_size", Value: "3"},
{Key: "hash_function", Value: "xxhash64"},
{Key: "seed", Value: "42"},
},
})
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorContains(t, err, "function output field not found")
})
t.Run("PreExecute happy path", func(t *testing.T) {
task := buildTask(buildValidRequest(), oldSchema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("PreExecute rejects invalid bound index name", func(t *testing.T) {
req := buildValidRequest()
req.GetAction().GetAddRequest().GetFieldInfos()[0].IndexName = "1bad-index-name"
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
assert.ErrorContains(t, err, "Invalid index name")
})
t.Run("PreExecute keeps the user's original index params in the request", func(t *testing.T) {
req := buildValidRequest()
original := proto.Clone(req.GetAction().GetAddRequest().GetFieldInfos()[0]).(*milvuspb.AlterCollectionSchemaRequest_FieldInfo)
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
// Validation-only: no normalization write-back, so the persisted
// UserIndexParams stay aligned with the create_index convention and a
// later create_index with identical params remains idempotent.
assert.True(t, proto.Equal(original, req.GetAction().GetAddRequest().GetFieldInfos()[0]))
})
t.Run("Execute leaves function output marker to RootCoord", func(t *testing.T) {
req := buildValidRequest()
task := buildTask(req, oldSchema)
// Ensure PreExecute passes first.
err := task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
addReq := req.GetAction().GetAddRequest()
assert.NotNil(t, addReq)
for _, fi := range addReq.GetFieldInfos() {
assert.False(t, fi.GetFieldSchema().GetIsFunctionOutput())
}
})
t.Run("Execute keeps plain added field as non-function output", func(t *testing.T) {
req := buildAddFieldRequest("execute_plain_text", true, true)
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
addReq := req.GetAction().GetAddRequest()
assert.NotNil(t, addReq)
for _, fi := range addReq.GetFieldInfos() {
assert.False(t, fi.GetFieldSchema().GetIsFunctionOutput())
}
})
t.Run("Execute supports function-only add request", func(t *testing.T) {
req := buildAddFunctionRequest(minHashFunctionOnlySchema)
task := buildTask(req, functionOnlyOldSchema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.NoError(t, err)
assert.Empty(t, req.GetAction().GetAddRequest().GetFieldInfos())
})
t.Run("Execute skips nil field infos", func(t *testing.T) {
req := buildValidRequest()
req.GetAction().GetAddRequest().FieldInfos = []*milvuspb.AlterCollectionSchemaRequest_FieldInfo{
nil,
{FieldSchema: nil},
}
task := buildTask(req, oldSchema)
err := task.Execute(ctx)
assert.NoError(t, err)
})
t.Run("Execute returns error when mixCoord AlterCollectionSchema fails", func(t *testing.T) {
mockey.PatchConvey("mixCoord returns error status", t, func() {
mockey.Mock((*MixCoordMock).AlterCollectionSchema).Return(
&milvuspb.AlterCollectionSchemaResponse{
AlterStatus: merr.Status(merr.ErrServiceUnavailable),
}, nil,
).Build()
req := buildValidRequest()
task := buildTask(req, oldSchema)
err := task.PreExecute(ctx)
assert.NoError(t, err)
err = task.Execute(ctx)
assert.Error(t, err)
})
})
}
// TestAlterCollection_RejectExternalTupleMutation verifies the proxy-layer
// guard rejects user alter_collection_properties calls that target the
// external_source / external_spec tuple. Mutation of this tuple is reserved
// to RefreshExternalCollection; see issue #49335.
func TestAlterCollection_RejectExternalTupleMutation(t *testing.T) {
ctx := context.Background()
task := &alterCollectionTask{
AlterCollectionRequest: &milvuspb.AlterCollectionRequest{
Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_AlterCollection},
DbName: dbName,
CollectionName: "ext_guard_test",
},
}
t.Run("reject properties with external_source", func(t *testing.T) {
task.Properties = []*commonpb.KeyValuePair{
{Key: common.CollectionExternalSource, Value: "s3://bucket/a/"},
}
task.DeleteKeys = nil
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "RefreshExternalCollection")
})
t.Run("reject properties with external_spec", func(t *testing.T) {
task.Properties = []*commonpb.KeyValuePair{
{Key: common.CollectionExternalSpec, Value: `{"format":"parquet"}`},
}
task.DeleteKeys = nil
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "RefreshExternalCollection")
})
t.Run("reject delete of external_source", func(t *testing.T) {
task.Properties = nil
task.DeleteKeys = []string{common.CollectionExternalSource}
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "immutable")
})
t.Run("reject delete of external_spec", func(t *testing.T) {
task.Properties = nil
task.DeleteKeys = []string{common.CollectionExternalSpec}
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "immutable")
})
}
func testDropFieldBaseSchema() *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vector", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
{FieldID: 102, Name: "scalar_field", DataType: schemapb.DataType_VarChar},
{FieldID: 103, Name: "partition_key", DataType: schemapb.DataType_Int64, IsPartitionKey: true},
{FieldID: 104, Name: "clustering_key", DataType: schemapb.DataType_Int64, IsClusteringKey: true},
{FieldID: 105, Name: "dynamic_field", DataType: schemapb.DataType_JSON, IsDynamic: true},
},
Functions: []*schemapb.FunctionSchema{
{Name: "test_func", InputFieldNames: []string{"scalar_field"}, OutputFieldNames: []string{"vector"}},
},
}
}
func TestValidateDropField(t *testing.T) {
baseSchema := testDropFieldBaseSchema()
t.Run("empty field name", func(t *testing.T) {
err := validateDropField(baseSchema, "")
assert.Error(t, err)
assert.Contains(t, err.Error(), "field name is empty")
})
t.Run("field not found", func(t *testing.T) {
err := validateDropField(baseSchema, "nonexistent_field")
assert.Error(t, err)
assert.Contains(t, err.Error(), "field not found")
})
t.Run("cannot drop primary key", func(t *testing.T) {
err := validateDropField(baseSchema, "id")
assert.Error(t, err)
assert.Contains(t, err.Error(), "cannot drop primary key field")
})
t.Run("cannot drop the last vector field", func(t *testing.T) {
err := validateDropField(baseSchema, "vector")
assert.Error(t, err)
assert.Contains(t, err.Error(), "cannot drop the last vector field")
})
t.Run("can drop vector field when multiple exist", func(t *testing.T) {
multiVecSchema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vector1", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
{FieldID: 102, Name: "vector2", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "64"}}},
},
}
err := validateDropField(multiVecSchema, "vector2")
assert.NoError(t, err)
})
t.Run("cannot drop partition key", func(t *testing.T) {
err := validateDropField(baseSchema, "partition_key")
assert.Error(t, err)
assert.Contains(t, err.Error(), "cannot drop partition key field")
})
t.Run("cannot drop clustering key", func(t *testing.T) {
err := validateDropField(baseSchema, "clustering_key")
assert.Error(t, err)
assert.Contains(t, err.Error(), "cannot drop clustering key field")
})
t.Run("cannot drop dynamic field", func(t *testing.T) {
err := validateDropField(baseSchema, "dynamic_field")
assert.Error(t, err)
assert.Contains(t, err.Error(), "cannot drop dynamic field")
})
t.Run("cannot drop active ttl field", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vector", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
{FieldID: 102, Name: "ttl", DataType: schemapb.DataType_Timestamptz, Nullable: true},
},
Properties: []*commonpb.KeyValuePair{
{Key: common.CollectionTTLFieldKey, Value: "ttl"},
},
}
err := validateDropField(schema, "ttl")
assert.Error(t, err)
assert.Contains(t, err.Error(), "referenced by collection property ttl_field")
})
t.Run("field referenced by function - input", func(t *testing.T) {
err := validateDropField(baseSchema, "scalar_field")
assert.Error(t, err)
assert.Contains(t, err.Error(), "is referenced by function test_func as input")
})
t.Run("field referenced by function - output", func(t *testing.T) {
// Need two vector fields so the "last vector field" check passes,
// allowing us to reach the output field reference check.
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec_output", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
{FieldID: 102, Name: "vec_other", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "64"}}},
{FieldID: 103, Name: "text", DataType: schemapb.DataType_VarChar},
},
Functions: []*schemapb.FunctionSchema{
{Name: "embed_func", InputFieldNames: []string{"text"}, OutputFieldNames: []string{"vec_output"}},
},
}
err := validateDropField(schema, "vec_output")
assert.Error(t, err)
assert.Contains(t, err.Error(), "is referenced by function embed_func as output")
})
t.Run("success - field can be dropped", func(t *testing.T) {
schemaNoFunc := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vector", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
{FieldID: 102, Name: "droppable_field", DataType: schemapb.DataType_VarChar},
},
}
err := validateDropField(schemaNoFunc, "droppable_field")
assert.NoError(t, err)
})
t.Run("can drop vector when struct array field also has vector", func(t *testing.T) {
// Top-level vector is the only Fields-level vector, but the struct
// array field also contains a vector, so total vector count == 2 and
// the drop should succeed.
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 102, Name: "paragraphs",
Fields: []*schemapb.FieldSchema{
{FieldID: 103, Name: "para_embed", DataType: schemapb.DataType_ArrayOfVector},
},
},
},
}
err := validateDropField(schema, "vec")
assert.NoError(t, err)
})
t.Run("drop whole struct array field by name - success", func(t *testing.T) {
// Top-level vector remains after dropping the struct, so the drop is allowed.
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 102, Name: "paragraphs",
Fields: []*schemapb.FieldSchema{
{FieldID: 103, Name: "para_text", DataType: schemapb.DataType_Array},
},
},
},
}
err := validateDropField(schema, "paragraphs")
assert.NoError(t, err)
})
t.Run("reject drop struct array field when it holds the only vector", func(t *testing.T) {
// The struct contains the only vector in the collection (an ArrayOfVector
// sub-field); dropping it would leave no vector field.
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 101, Name: "paragraphs",
Fields: []*schemapb.FieldSchema{
{FieldID: 102, Name: "para_text", DataType: schemapb.DataType_Array},
{FieldID: 103, Name: "para_embed", DataType: schemapb.DataType_ArrayOfVector},
},
},
},
}
err := validateDropField(schema, "paragraphs")
assert.Error(t, err)
assert.Contains(t, err.Error(), "would leave no vector field")
})
t.Run("reject drop struct array field with protected sub-field", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 102, Name: "paragraphs",
Fields: []*schemapb.FieldSchema{
{FieldID: 103, Name: "para_key", DataType: schemapb.DataType_Array, IsPartitionKey: true},
},
},
},
}
err := validateDropField(schema, "paragraphs")
assert.Error(t, err)
assert.Contains(t, err.Error(), "sub-field para_key is partition key")
})
t.Run("reject drop struct array field referenced by function", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 102, Name: "paragraphs",
Fields: []*schemapb.FieldSchema{
{FieldID: 103, Name: "para_text", DataType: schemapb.DataType_Array},
},
},
},
Functions: []*schemapb.FunctionSchema{
{Name: "embed_func", InputFieldNames: []string{"para_text"}, OutputFieldNames: []string{"vec"}},
},
}
err := validateDropField(schema, "paragraphs")
assert.Error(t, err)
assert.Contains(t, err.Error(), "sub-field para_text is referenced by function embed_func as input")
})
t.Run("reject drop sub-field of struct array field by name", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 102, Name: "paragraphs",
Fields: []*schemapb.FieldSchema{
{FieldID: 103, Name: "para_text", DataType: schemapb.DataType_Array},
},
},
},
}
err := validateDropField(schema, "para_text")
assert.Error(t, err)
assert.Contains(t, err.Error(), "cannot drop sub-field of struct array field")
assert.Contains(t, err.Error(), "paragraphs.para_text")
})
}
func TestAlterCollectionSchemaTask_PreExecute(t *testing.T) {
ctx := context.Background()
baseSchema := testDropFieldBaseSchema()
t.Run("nil action", func(t *testing.T) {
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: baseSchema,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: nil,
},
}
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "action is nil")
})
t.Run("empty add request fails validation", func(t *testing.T) {
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: baseSchema,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_AddRequest{},
},
},
}
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "add_request is nil")
})
t.Run("drop by field_name - validation error", func(t *testing.T) {
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: baseSchema,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_DropRequest{
DropRequest: &milvuspb.AlterCollectionSchemaRequest_DropRequest{
Identifier: &milvuspb.AlterCollectionSchemaRequest_DropRequest_FieldName{
FieldName: "id",
},
},
},
},
},
}
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "cannot drop primary key field")
})
t.Run("drop by field_id - success", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vector", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
{FieldID: 102, Name: "droppable", DataType: schemapb.DataType_VarChar},
},
}
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: schema,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_DropRequest{
DropRequest: &milvuspb.AlterCollectionSchemaRequest_DropRequest{
Identifier: &milvuspb.AlterCollectionSchemaRequest_DropRequest_FieldId{
FieldId: 102,
},
},
},
},
},
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("drop by field_id - not found", func(t *testing.T) {
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: baseSchema,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_DropRequest{
DropRequest: &milvuspb.AlterCollectionSchemaRequest_DropRequest{
Identifier: &milvuspb.AlterCollectionSchemaRequest_DropRequest_FieldId{
FieldId: 999,
},
},
},
},
},
}
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "field not found with id")
})
t.Run("drop by field_name - success", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vector", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
{FieldID: 102, Name: "droppable", DataType: schemapb.DataType_VarChar},
},
}
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: schema,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_DropRequest{
DropRequest: &milvuspb.AlterCollectionSchemaRequest_DropRequest{
Identifier: &milvuspb.AlterCollectionSchemaRequest_DropRequest_FieldName{
FieldName: "droppable",
},
},
},
},
},
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("drop by function_name - detach minhash success", func(t *testing.T) {
schemaWithFunc := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vector", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
{FieldID: 102, Name: "text", DataType: schemapb.DataType_VarChar},
},
Functions: []*schemapb.FunctionSchema{
{Name: "minhash_func", Type: schemapb.FunctionType_MinHash},
},
}
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: schemaWithFunc,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_DropRequest{
DropRequest: &milvuspb.AlterCollectionSchemaRequest_DropRequest{
Identifier: &milvuspb.AlterCollectionSchemaRequest_DropRequest_FunctionName{
FunctionName: "minhash_func",
},
},
},
},
},
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("drop by function_name - bm25 function field success", func(t *testing.T) {
schemaWithFunc := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "dense_vector", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
{FieldID: 102, Name: "text", DataType: schemapb.DataType_VarChar},
{FieldID: 103, Name: "sparse_vector", DataType: schemapb.DataType_SparseFloatVector},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "bm25_func",
Type: schemapb.FunctionType_BM25,
InputFieldNames: []string{"text"},
OutputFieldNames: []string{"sparse_vector"},
},
},
}
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: schemaWithFunc,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_DropRequest{
DropRequest: &milvuspb.AlterCollectionSchemaRequest_DropRequest{
Identifier: &milvuspb.AlterCollectionSchemaRequest_DropRequest_FunctionName{
FunctionName: "bm25_func",
},
DropFunctionOutputFields: true,
},
},
},
},
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("drop by function_name - not found", func(t *testing.T) {
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: baseSchema,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_DropRequest{
DropRequest: &milvuspb.AlterCollectionSchemaRequest_DropRequest{
Identifier: &milvuspb.AlterCollectionSchemaRequest_DropRequest_FunctionName{
FunctionName: "nonexistent_func",
},
},
},
},
},
}
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "function not found")
})
t.Run("drop with empty schema", func(t *testing.T) {
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: nil,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_DropRequest{
DropRequest: &milvuspb.AlterCollectionSchemaRequest_DropRequest{
Identifier: &milvuspb.AlterCollectionSchemaRequest_DropRequest_FieldName{
FieldName: "field",
},
},
},
},
},
}
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "empty old schema")
})
t.Run("drop by field_id - struct array field success", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 102, Name: "paragraphs",
Fields: []*schemapb.FieldSchema{
{FieldID: 103, Name: "para_text", DataType: schemapb.DataType_Array},
},
},
},
}
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: schema,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_DropRequest{
DropRequest: &milvuspb.AlterCollectionSchemaRequest_DropRequest{
Identifier: &milvuspb.AlterCollectionSchemaRequest_DropRequest_FieldId{
FieldId: 102,
},
},
},
},
},
}
err := task.PreExecute(ctx)
assert.NoError(t, err)
})
t.Run("drop by field_id - reject sub-field of struct array field", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}}},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 102, Name: "paragraphs",
Fields: []*schemapb.FieldSchema{
{FieldID: 103, Name: "para_text", DataType: schemapb.DataType_Array},
},
},
},
}
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: schema,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{
Op: &milvuspb.AlterCollectionSchemaRequest_Action_DropRequest{
DropRequest: &milvuspb.AlterCollectionSchemaRequest_DropRequest{
Identifier: &milvuspb.AlterCollectionSchemaRequest_DropRequest_FieldId{
FieldId: 103,
},
},
},
},
},
}
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "cannot drop sub-field of struct array field")
assert.Contains(t, err.Error(), "paragraphs.para_text")
})
t.Run("unknown action type", func(t *testing.T) {
task := &alterCollectionSchemaTask{
ctx: ctx,
oldSchema: baseSchema,
AlterCollectionSchemaRequest: &milvuspb.AlterCollectionSchemaRequest{
CollectionName: "test_collection",
Action: &milvuspb.AlterCollectionSchemaRequest_Action{},
},
}
err := task.PreExecute(ctx)
assert.Error(t, err)
assert.Contains(t, err.Error(), "unknown action type")
})
}
func TestValidateDropFunction(t *testing.T) {
t.Run("empty function name", func(t *testing.T) {
err := validateDropFunction(&schemapb.CollectionSchema{}, "", false)
assert.Error(t, err)
assert.Contains(t, err.Error(), "function name is empty")
})
t.Run("function not found", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Functions: []*schemapb.FunctionSchema{
{Name: "embedding_func", Type: schemapb.FunctionType_TextEmbedding},
},
}
err := validateDropFunction(schema, "nonexistent", false)
assert.Error(t, err)
assert.Contains(t, err.Error(), "function not found")
})
t.Run("detach minhash function succeeds", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
{FieldID: 102, Name: "minhash_vec", DataType: schemapb.DataType_BinaryVector},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "minhash_func", Type: schemapb.FunctionType_MinHash,
InputFieldNames: []string{"text"}, OutputFieldNames: []string{"minhash_vec"},
},
},
}
err := validateDropFunction(schema, "minhash_func", false)
assert.NoError(t, err)
})
t.Run("detach bm25 function fails", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
{FieldID: 102, Name: "sparse_vec", DataType: schemapb.DataType_SparseFloatVector},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "bm25_func", Type: schemapb.FunctionType_BM25,
InputFieldNames: []string{"text"}, OutputFieldNames: []string{"sparse_vec"},
},
},
}
err := validateDropFunction(schema, "bm25_func", false)
assert.Error(t, err)
assert.Contains(t, err.Error(), "BM25 function must be dropped with its output field")
})
t.Run("drop function field leaves another vector field", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
{FieldID: 102, Name: "sparse_vec", DataType: schemapb.DataType_SparseFloatVector},
{FieldID: 103, Name: "vec_other", DataType: schemapb.DataType_FloatVector},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "bm25_func", Type: schemapb.FunctionType_BM25,
InputFieldNames: []string{"text"}, OutputFieldNames: []string{"sparse_vec"},
},
},
}
err := validateDropFunction(schema, "bm25_func", true)
assert.NoError(t, err)
})
t.Run("drop minhash function field succeeds", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
{FieldID: 102, Name: "minhash_vec", DataType: schemapb.DataType_BinaryVector},
{FieldID: 103, Name: "vec_other", DataType: schemapb.DataType_FloatVector},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "minhash_func", Type: schemapb.FunctionType_MinHash,
InputFieldNames: []string{"text"}, OutputFieldNames: []string{"minhash_vec"},
},
},
}
err := validateDropFunction(schema, "minhash_func", true)
assert.NoError(t, err)
})
t.Run("drop unsupported function field fails", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
{FieldID: 102, Name: "vec_func_out", DataType: schemapb.DataType_FloatVector},
{FieldID: 103, Name: "vec_other", DataType: schemapb.DataType_FloatVector},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "embed_func", Type: schemapb.FunctionType_TextEmbedding,
InputFieldNames: []string{"text"}, OutputFieldNames: []string{"vec_func_out"},
},
},
}
err := validateDropFunction(schema, "embed_func", true)
assert.Error(t, err)
assert.Contains(t, err.Error(), "only BM25 and MinHash functions support dropping output fields")
})
t.Run("drop function field would leave no vector field", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
{FieldID: 102, Name: "sparse_vec", DataType: schemapb.DataType_SparseFloatVector},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "bm25_func", Type: schemapb.FunctionType_BM25,
InputFieldNames: []string{"text"}, OutputFieldNames: []string{"sparse_vec"},
},
},
}
err := validateDropFunction(schema, "bm25_func", true)
assert.Error(t, err)
assert.Contains(t, err.Error(), "would leave no vector field")
})
t.Run("drop function field preserves vector in struct array field", func(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
{FieldID: 102, Name: "sparse_vec", DataType: schemapb.DataType_SparseFloatVector},
},
StructArrayFields: []*schemapb.StructArrayFieldSchema{
{
FieldID: 103, Name: "paragraphs",
Fields: []*schemapb.FieldSchema{
{FieldID: 104, Name: "para_embed", DataType: schemapb.DataType_ArrayOfVector},
},
},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "bm25_func", Type: schemapb.FunctionType_BM25,
InputFieldNames: []string{"text"}, OutputFieldNames: []string{"sparse_vec"},
},
},
}
err := validateDropFunction(schema, "bm25_func", true)
assert.NoError(t, err)
})
}