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
815 lines
24 KiB
Go
815 lines
24 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 task
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
|
"github.com/milvus-io/milvus/pkg/v3/common"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
|
)
|
|
|
|
type UtilsSuite struct {
|
|
suite.Suite
|
|
}
|
|
|
|
func (s *UtilsSuite) TestPackLoadMetaSchemaVersions() {
|
|
collectionInfoResp := &milvuspb.DescribeCollectionResponse{
|
|
CollectionID: 100,
|
|
DbName: "default",
|
|
UpdateTimestamp: 200,
|
|
Schema: &schemapb.CollectionSchema{
|
|
Version: 3,
|
|
},
|
|
}
|
|
|
|
loadMeta := packLoadMeta(querypb.LoadType_LoadCollection, collectionInfoResp, "rg", []int64{10}, 20)
|
|
s.Equal(uint64(200), loadMeta.GetSchemaBarrierTs())
|
|
}
|
|
|
|
func (s *UtilsSuite) TestPackLoadSegmentRequest() {
|
|
ctx := context.Background()
|
|
|
|
action := NewSegmentAction(1, ActionTypeGrow, "test-ch", 100)
|
|
task, err := NewSegmentTask(
|
|
ctx,
|
|
time.Second,
|
|
nil,
|
|
1,
|
|
newReplicaDefaultRG(10),
|
|
commonpb.LoadPriority_LOW,
|
|
action,
|
|
)
|
|
s.NoError(err)
|
|
|
|
collectionInfoResp := &milvuspb.DescribeCollectionResponse{
|
|
Schema: &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
},
|
|
},
|
|
},
|
|
Properties: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.MmapEnabledKey,
|
|
Value: "false",
|
|
},
|
|
},
|
|
}
|
|
|
|
req := packLoadSegmentRequest(
|
|
task,
|
|
action,
|
|
collectionInfoResp.GetSchema(),
|
|
collectionInfoResp.GetProperties(),
|
|
&querypb.LoadMetaInfo{
|
|
LoadType: querypb.LoadType_LoadCollection,
|
|
},
|
|
&querypb.SegmentLoadInfo{},
|
|
nil,
|
|
)
|
|
|
|
s.True(req.GetNeedTransfer())
|
|
s.Equal(task.CollectionID(), req.CollectionID)
|
|
s.Equal(task.ReplicaID(), req.ReplicaID)
|
|
s.Equal(action.Node(), req.GetDstNodeID())
|
|
for _, field := range req.GetSchema().GetFields() {
|
|
mmapEnable, ok := common.IsMmapDataEnabled(field.GetTypeParams()...)
|
|
s.False(mmapEnable)
|
|
s.True(ok)
|
|
}
|
|
}
|
|
|
|
func (s *UtilsSuite) TestPackLoadSegmentRequestMmapDuplicateBug() {
|
|
// This test demonstrates the bug where mmap.enabled is appended twice to field TypeParams
|
|
ctx := context.Background()
|
|
|
|
action := NewSegmentAction(1, ActionTypeGrow, "test-ch", 100)
|
|
task, err := NewSegmentTask(
|
|
ctx,
|
|
time.Second,
|
|
nil,
|
|
1,
|
|
newReplicaDefaultRG(10),
|
|
commonpb.LoadPriority_LOW,
|
|
action,
|
|
)
|
|
s.NoError(err)
|
|
|
|
collectionInfoResp := &milvuspb.DescribeCollectionResponse{
|
|
Schema: &schemapb.CollectionSchema{
|
|
Name: "test_collection",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "pk",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
TypeParams: []*commonpb.KeyValuePair{}, // No existing mmap setting
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "vector",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "dim", Value: "128"},
|
|
}, // No existing mmap setting
|
|
},
|
|
{
|
|
FieldID: 102,
|
|
Name: "scalar",
|
|
DataType: schemapb.DataType_Int32,
|
|
TypeParams: []*commonpb.KeyValuePair{}, // No existing mmap setting
|
|
},
|
|
},
|
|
},
|
|
Properties: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.MmapEnabledKey,
|
|
Value: "true",
|
|
},
|
|
},
|
|
}
|
|
|
|
req := packLoadSegmentRequest(
|
|
task,
|
|
action,
|
|
collectionInfoResp.GetSchema(),
|
|
collectionInfoResp.GetProperties(),
|
|
&querypb.LoadMetaInfo{
|
|
LoadType: querypb.LoadType_LoadCollection,
|
|
},
|
|
&querypb.SegmentLoadInfo{},
|
|
nil,
|
|
)
|
|
|
|
// Check each field for duplicate mmap.enabled properties
|
|
for _, field := range req.GetSchema().GetFields() {
|
|
mmapCount := 0
|
|
for _, kv := range field.GetTypeParams() {
|
|
if kv.Key == common.MmapEnabledKey {
|
|
mmapCount++
|
|
}
|
|
}
|
|
s.LessOrEqualf(mmapCount, 1, "Field %s has %d mmap.enabled properties, should have at most 1", field.GetName(), mmapCount)
|
|
}
|
|
}
|
|
|
|
func (s *UtilsSuite) TestPackLoadSegmentRequestMmap() {
|
|
ctx := context.Background()
|
|
|
|
action := NewSegmentAction(1, ActionTypeGrow, "test-ch", 100)
|
|
task, err := NewSegmentTask(
|
|
ctx,
|
|
time.Second,
|
|
nil,
|
|
1,
|
|
newReplicaDefaultRG(10),
|
|
commonpb.LoadPriority_LOW,
|
|
action,
|
|
)
|
|
s.NoError(err)
|
|
|
|
collectionInfoResp := &milvuspb.DescribeCollectionResponse{
|
|
Schema: &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
},
|
|
},
|
|
},
|
|
Properties: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.MmapEnabledKey,
|
|
Value: "true",
|
|
},
|
|
},
|
|
}
|
|
|
|
req := packLoadSegmentRequest(
|
|
task,
|
|
action,
|
|
collectionInfoResp.GetSchema(),
|
|
collectionInfoResp.GetProperties(),
|
|
&querypb.LoadMetaInfo{
|
|
LoadType: querypb.LoadType_LoadCollection,
|
|
},
|
|
&querypb.SegmentLoadInfo{},
|
|
nil,
|
|
)
|
|
|
|
s.True(req.GetNeedTransfer())
|
|
s.Equal(task.CollectionID(), req.CollectionID)
|
|
s.Equal(task.ReplicaID(), req.ReplicaID)
|
|
s.Equal(action.Node(), req.GetDstNodeID())
|
|
for _, field := range req.GetSchema().GetFields() {
|
|
mmapEnable, ok := common.IsMmapDataEnabled(field.GetTypeParams()...)
|
|
s.True(mmapEnable)
|
|
s.True(ok)
|
|
}
|
|
}
|
|
|
|
func (s *UtilsSuite) TestPackLoadSegmentRequestFieldLevelMmapPriority() {
|
|
// This test demonstrates that field-level mmap settings should not be duplicated
|
|
// when collection-level mmap is also set
|
|
ctx := context.Background()
|
|
|
|
action := NewSegmentAction(1, ActionTypeGrow, "test-ch", 100)
|
|
task, err := NewSegmentTask(
|
|
ctx,
|
|
time.Second,
|
|
nil,
|
|
1,
|
|
newReplicaDefaultRG(10),
|
|
commonpb.LoadPriority_LOW,
|
|
action,
|
|
)
|
|
s.NoError(err)
|
|
|
|
collectionInfoResp := &milvuspb.DescribeCollectionResponse{
|
|
Schema: &schemapb.CollectionSchema{
|
|
Name: "test_collection",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "pk",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
TypeParams: []*commonpb.KeyValuePair{}, // No field-level mmap
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "vector_field_mmap_disabled",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "dim", Value: "128"},
|
|
{Key: common.MmapEnabledKey, Value: "false"}, // Field-level mmap explicitly disabled
|
|
},
|
|
},
|
|
{
|
|
FieldID: 102,
|
|
Name: "vector_field_mmap_enabled",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "dim", Value: "256"},
|
|
{Key: common.MmapEnabledKey, Value: "true"}, // Field-level mmap explicitly enabled
|
|
},
|
|
},
|
|
},
|
|
},
|
|
Properties: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.MmapEnabledKey,
|
|
Value: "true", // Collection-level mmap enabled
|
|
},
|
|
},
|
|
}
|
|
|
|
req := packLoadSegmentRequest(
|
|
task,
|
|
action,
|
|
collectionInfoResp.GetSchema(),
|
|
collectionInfoResp.GetProperties(),
|
|
&querypb.LoadMetaInfo{
|
|
LoadType: querypb.LoadType_LoadCollection,
|
|
},
|
|
&querypb.SegmentLoadInfo{},
|
|
nil,
|
|
)
|
|
|
|
// Check field without field-level mmap setting
|
|
pkField := req.GetSchema().GetFields()[0]
|
|
mmapCount := 0
|
|
var mmapValue string
|
|
for _, kv := range pkField.GetTypeParams() {
|
|
if kv.Key == common.MmapEnabledKey {
|
|
mmapCount++
|
|
mmapValue = kv.Value
|
|
}
|
|
}
|
|
s.Equalf(1, mmapCount, "Field %s should have exactly 1 mmap.enabled property", pkField.GetName())
|
|
s.Equalf("true", mmapValue, "Field %s should inherit collection-level mmap=true", pkField.GetName())
|
|
|
|
// Check field with field-level mmap=false (should not be overridden by collection-level)
|
|
vectorField1 := req.GetSchema().GetFields()[1]
|
|
mmapCount = 0
|
|
mmapValues := []string{}
|
|
for _, kv := range vectorField1.GetTypeParams() {
|
|
if kv.Key == common.MmapEnabledKey {
|
|
mmapCount++
|
|
mmapValues = append(mmapValues, kv.Value)
|
|
}
|
|
}
|
|
|
|
s.Equalf(1, mmapCount, "Field %s should have exactly 1 mmap.enabled property, got %d with values %v",
|
|
vectorField1.GetName(), mmapCount, mmapValues)
|
|
if mmapCount > 0 {
|
|
s.Equalf("false", mmapValues[0], "Field-level mmap=false should be preserved for field %s", vectorField1.GetName())
|
|
}
|
|
|
|
// Check field with field-level mmap=true (should not be duplicated)
|
|
vectorField2 := req.GetSchema().GetFields()[2]
|
|
mmapCount = 0
|
|
mmapValues = []string{}
|
|
for _, kv := range vectorField2.GetTypeParams() {
|
|
if kv.Key == common.MmapEnabledKey {
|
|
mmapCount++
|
|
mmapValues = append(mmapValues, kv.Value)
|
|
}
|
|
}
|
|
s.Equalf(1, mmapCount, "Field %s should have exactly 1 mmap.enabled property, got %d with values %v",
|
|
vectorField2.GetName(), mmapCount, mmapValues)
|
|
}
|
|
|
|
// TestFieldMmapPriorityOverCollection tests that field-level mmap settings have priority over collection-level settings
|
|
// The priority rule should be: field-level mmap > collection-level mmap > default (no mmap)
|
|
func (s *UtilsSuite) TestFieldMmapPriorityOverCollection() {
|
|
ctx := context.Background()
|
|
|
|
action := NewSegmentAction(1, ActionTypeGrow, "test-ch", 100)
|
|
task, err := NewSegmentTask(
|
|
ctx,
|
|
time.Second,
|
|
nil,
|
|
1,
|
|
newReplicaDefaultRG(10),
|
|
commonpb.LoadPriority_LOW,
|
|
action,
|
|
)
|
|
s.NoError(err)
|
|
|
|
// Test case 1: Collection mmap=true, field mmap=false -> field should remain false
|
|
s.Run("collection_true_field_false", func() {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test_collection",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "field_with_mmap_false",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "dim", Value: "128"},
|
|
{Key: common.MmapEnabledKey, Value: "false"}, // Field explicitly sets mmap=false
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
collectionProps := []*commonpb.KeyValuePair{
|
|
{Key: common.MmapEnabledKey, Value: "true"}, // Collection sets mmap=true
|
|
}
|
|
|
|
req := packLoadSegmentRequest(
|
|
task,
|
|
action,
|
|
schema,
|
|
collectionProps,
|
|
&querypb.LoadMetaInfo{LoadType: querypb.LoadType_LoadCollection},
|
|
&querypb.SegmentLoadInfo{},
|
|
nil,
|
|
)
|
|
|
|
field := req.GetSchema().GetFields()[0]
|
|
mmapCount := 0
|
|
mmapValues := []string{}
|
|
for _, kv := range field.GetTypeParams() {
|
|
if kv.Key == common.MmapEnabledKey {
|
|
mmapCount++
|
|
mmapValues = append(mmapValues, kv.Value)
|
|
}
|
|
}
|
|
|
|
// Should have exactly 1 mmap property
|
|
s.Equalf(1, mmapCount, "Field should have exactly 1 mmap.enabled, got %d with values %v", mmapCount, mmapValues)
|
|
// Field-level false should override collection-level true
|
|
if mmapCount > 0 {
|
|
s.Equal("false", mmapValues[0], "Field mmap=false should have priority over collection mmap=true")
|
|
}
|
|
})
|
|
|
|
// Test case 2: Collection mmap=false, field mmap=true -> field should remain true
|
|
s.Run("collection_false_field_true", func() {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test_collection",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "field_with_mmap_true",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "dim", Value: "128"},
|
|
{Key: common.MmapEnabledKey, Value: "true"}, // Field explicitly sets mmap=true
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
collectionProps := []*commonpb.KeyValuePair{
|
|
{Key: common.MmapEnabledKey, Value: "false"}, // Collection sets mmap=false
|
|
}
|
|
|
|
req := packLoadSegmentRequest(
|
|
task,
|
|
action,
|
|
schema,
|
|
collectionProps,
|
|
&querypb.LoadMetaInfo{LoadType: querypb.LoadType_LoadCollection},
|
|
&querypb.SegmentLoadInfo{},
|
|
nil,
|
|
)
|
|
|
|
field := req.GetSchema().GetFields()[0]
|
|
mmapCount := 0
|
|
mmapValues := []string{}
|
|
for _, kv := range field.GetTypeParams() {
|
|
if kv.Key == common.MmapEnabledKey {
|
|
mmapCount++
|
|
mmapValues = append(mmapValues, kv.Value)
|
|
}
|
|
}
|
|
|
|
// Should have exactly 1 mmap property
|
|
s.Equalf(1, mmapCount, "Field should have exactly 1 mmap.enabled, got %d with values %v", mmapCount, mmapValues)
|
|
// Field-level true should override collection-level false
|
|
if mmapCount > 0 {
|
|
s.Equal("true", mmapValues[0], "Field mmap=true should have priority over collection mmap=false")
|
|
}
|
|
})
|
|
|
|
// Test case 3: Mixed fields - some with field-level settings, some without
|
|
s.Run("mixed_fields", func() {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test_collection",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "no_mmap_setting",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
TypeParams: []*commonpb.KeyValuePair{}, // No field-level mmap
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "field_mmap_false",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "dim", Value: "128"},
|
|
{Key: common.MmapEnabledKey, Value: "false"},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 102,
|
|
Name: "field_mmap_true",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "dim", Value: "256"},
|
|
{Key: common.MmapEnabledKey, Value: "true"},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 103,
|
|
Name: "no_mmap_scalar",
|
|
DataType: schemapb.DataType_VarChar,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: "max_length", Value: "100"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
collectionProps := []*commonpb.KeyValuePair{
|
|
{Key: common.MmapEnabledKey, Value: "true"}, // Collection-level mmap=true
|
|
}
|
|
|
|
req := packLoadSegmentRequest(
|
|
task,
|
|
action,
|
|
schema,
|
|
collectionProps,
|
|
&querypb.LoadMetaInfo{LoadType: querypb.LoadType_LoadCollection},
|
|
&querypb.SegmentLoadInfo{},
|
|
nil,
|
|
)
|
|
|
|
fields := req.GetSchema().GetFields()
|
|
s.Equal(4, len(fields))
|
|
|
|
// Check each field
|
|
for i, expectedValues := range []struct {
|
|
fieldName string
|
|
expectedMmap string
|
|
description string
|
|
}{
|
|
{"no_mmap_setting", "true", "Field without setting should inherit collection mmap=true"},
|
|
{"field_mmap_false", "false", "Field with mmap=false should keep false despite collection mmap=true"},
|
|
{"field_mmap_true", "true", "Field with mmap=true should keep true"},
|
|
{"no_mmap_scalar", "true", "Scalar field without setting should inherit collection mmap=true"},
|
|
} {
|
|
field := fields[i]
|
|
s.Equal(expectedValues.fieldName, field.GetName())
|
|
|
|
mmapCount := 0
|
|
var actualMmap string
|
|
for _, kv := range field.GetTypeParams() {
|
|
if kv.Key == common.MmapEnabledKey {
|
|
mmapCount++
|
|
actualMmap = kv.Value
|
|
}
|
|
}
|
|
|
|
// Check no duplicates
|
|
s.Equalf(1, mmapCount, "Field %s should have exactly 1 mmap.enabled property, got %d",
|
|
field.GetName(), mmapCount)
|
|
|
|
// Check correct value based on priority
|
|
s.Equalf(expectedValues.expectedMmap, actualMmap,
|
|
"%s for field %s", expectedValues.description, field.GetName())
|
|
}
|
|
})
|
|
}
|
|
|
|
func (s *UtilsSuite) TestPackLoadSegmentRequestLoadScope() {
|
|
ctx := context.Background()
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
},
|
|
},
|
|
}
|
|
|
|
testCases := []struct {
|
|
name string
|
|
action ActionType
|
|
wantScope querypb.LoadScope
|
|
}{
|
|
{
|
|
name: "index update",
|
|
action: ActionTypeUpdate,
|
|
wantScope: querypb.LoadScope_Index,
|
|
},
|
|
{
|
|
name: "legacy stats update",
|
|
action: ActionTypeStatsUpdate,
|
|
wantScope: querypb.LoadScope_Stats,
|
|
},
|
|
{
|
|
name: "reopen update",
|
|
action: ActionTypeReopen,
|
|
wantScope: querypb.LoadScope_Reopen,
|
|
},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
s.Run(tc.name, func() {
|
|
action := NewSegmentAction(1, tc.action, "test-ch", 100)
|
|
task, err := NewSegmentTask(
|
|
ctx,
|
|
time.Second,
|
|
nil,
|
|
1,
|
|
newReplicaDefaultRG(10),
|
|
commonpb.LoadPriority_LOW,
|
|
action,
|
|
)
|
|
s.NoError(err)
|
|
|
|
req := packLoadSegmentRequest(
|
|
task,
|
|
action,
|
|
schema,
|
|
nil,
|
|
&querypb.LoadMetaInfo{LoadType: querypb.LoadType_LoadCollection},
|
|
&querypb.SegmentLoadInfo{},
|
|
nil,
|
|
)
|
|
|
|
s.Equal(tc.wantScope, req.GetLoadScope())
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAutoWarmupForNonPKIsolationCollection(t *testing.T) {
|
|
paramtable.Init()
|
|
|
|
t.Run("config disabled", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key, "false")
|
|
defer paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key)
|
|
result := autoWarmupForNonPKIsolationCollection(nil)
|
|
assert.False(t, result)
|
|
})
|
|
|
|
t.Run("config enabled, no PKI", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key)
|
|
result := autoWarmupForNonPKIsolationCollection(nil)
|
|
assert.True(t, result)
|
|
})
|
|
|
|
t.Run("config enabled, with PKI", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key)
|
|
result := autoWarmupForNonPKIsolationCollection([]*commonpb.KeyValuePair{
|
|
{Key: common.PartitionKeyIsolationKey, Value: "true"},
|
|
})
|
|
assert.False(t, result)
|
|
})
|
|
|
|
t.Run("config enabled, PKI explicitly false", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key)
|
|
result := autoWarmupForNonPKIsolationCollection([]*commonpb.KeyValuePair{
|
|
{Key: common.PartitionKeyIsolationKey, Value: "false"},
|
|
})
|
|
assert.True(t, result)
|
|
})
|
|
}
|
|
|
|
func TestApplyCollectionWarmupSettingAutoWarmup(t *testing.T) {
|
|
paramtable.Init()
|
|
|
|
t.Run("autoWarmupForNonPKIsolationCollection sets sync for scalar fields, skips vector fields", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key)
|
|
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 1, DataType: schemapb.DataType_Int64},
|
|
{FieldID: 2, DataType: schemapb.DataType_VarChar},
|
|
{FieldID: 3, DataType: schemapb.DataType_FloatVector},
|
|
},
|
|
}
|
|
collectionProps := []*commonpb.KeyValuePair{}
|
|
|
|
result := applyCollectionWarmupSetting(schema, collectionProps)
|
|
|
|
// Scalar fields should get sync warmup
|
|
for _, field := range result.GetFields() {
|
|
warmup, exist := common.GetWarmupPolicy(field.GetTypeParams()...)
|
|
if field.GetDataType() == schemapb.DataType_FloatVector {
|
|
assert.False(t, exist, "vector field should NOT have warmup set by autoWarmupForNonPKIsolationCollection")
|
|
} else {
|
|
assert.True(t, exist, "scalar field should have warmup set by autoWarmupForNonPKIsolationCollection")
|
|
assert.Equal(t, common.WarmupSync, warmup)
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("autoWarmupForNonPKIsolationCollection disabled for PKI collection", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key)
|
|
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 1, DataType: schemapb.DataType_Int64},
|
|
},
|
|
}
|
|
collectionProps := []*commonpb.KeyValuePair{
|
|
{Key: common.PartitionKeyIsolationKey, Value: "true"},
|
|
}
|
|
|
|
result := applyCollectionWarmupSetting(schema, collectionProps)
|
|
|
|
field := result.GetFields()[0]
|
|
_, exist := common.GetWarmupPolicy(field.GetTypeParams()...)
|
|
assert.False(t, exist, "PKI collection should not have autoWarmupForNonPKIsolationCollection")
|
|
})
|
|
|
|
t.Run("collection-level warmup takes priority over autoWarmupForNonPKIsolationCollection", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key)
|
|
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 1, DataType: schemapb.DataType_Int64},
|
|
},
|
|
}
|
|
collectionProps := []*commonpb.KeyValuePair{
|
|
{Key: common.WarmupScalarFieldKey, Value: common.WarmupDisable},
|
|
}
|
|
|
|
result := applyCollectionWarmupSetting(schema, collectionProps)
|
|
|
|
field := result.GetFields()[0]
|
|
warmup, exist := common.GetWarmupPolicy(field.GetTypeParams()...)
|
|
assert.True(t, exist)
|
|
assert.Equal(t, common.WarmupDisable, warmup, "collection-level setting should override autoWarmupForNonPKIsolationCollection")
|
|
})
|
|
}
|
|
|
|
func TestApplyIndexWarmupSettingAutoWarmup(t *testing.T) {
|
|
paramtable.Init()
|
|
|
|
t.Run("autoWarmupForNonPKIsolationCollection sets sync for all indexes", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key)
|
|
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 1, DataType: schemapb.DataType_Int64},
|
|
{FieldID: 2, DataType: schemapb.DataType_FloatVector},
|
|
},
|
|
}
|
|
loadInfo := &querypb.SegmentLoadInfo{
|
|
IndexInfos: []*querypb.FieldIndexInfo{
|
|
{FieldID: 1, IndexParams: []*commonpb.KeyValuePair{}},
|
|
{FieldID: 2, IndexParams: []*commonpb.KeyValuePair{}},
|
|
},
|
|
}
|
|
collectionProps := []*commonpb.KeyValuePair{}
|
|
|
|
applyIndexWarmupSetting(loadInfo, schema, collectionProps)
|
|
|
|
for _, indexInfo := range loadInfo.GetIndexInfos() {
|
|
warmup, exist := common.GetWarmupPolicy(indexInfo.IndexParams...)
|
|
assert.True(t, exist, "all indexes should have warmup set by autoWarmupForNonPKIsolationCollection")
|
|
assert.Equal(t, common.WarmupSync, warmup)
|
|
}
|
|
})
|
|
|
|
t.Run("autoWarmupForNonPKIsolationCollection disabled for PKI collection indexes", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key)
|
|
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 1, DataType: schemapb.DataType_Int64},
|
|
},
|
|
}
|
|
loadInfo := &querypb.SegmentLoadInfo{
|
|
IndexInfos: []*querypb.FieldIndexInfo{
|
|
{FieldID: 1, IndexParams: []*commonpb.KeyValuePair{}},
|
|
},
|
|
}
|
|
collectionProps := []*commonpb.KeyValuePair{
|
|
{Key: common.PartitionKeyIsolationKey, Value: "true"},
|
|
}
|
|
|
|
applyIndexWarmupSetting(loadInfo, schema, collectionProps)
|
|
|
|
indexInfo := loadInfo.GetIndexInfos()[0]
|
|
_, exist := common.GetWarmupPolicy(indexInfo.IndexParams...)
|
|
assert.False(t, exist, "PKI collection index should not have autoWarmupForNonPKIsolationCollection")
|
|
})
|
|
|
|
t.Run("collection-level index warmup takes priority over autoWarmupForNonPKIsolationCollection", func(t *testing.T) {
|
|
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key, "true")
|
|
defer paramtable.Get().Reset(paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.Key)
|
|
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 1, DataType: schemapb.DataType_Int64},
|
|
},
|
|
}
|
|
loadInfo := &querypb.SegmentLoadInfo{
|
|
IndexInfos: []*querypb.FieldIndexInfo{
|
|
{FieldID: 1, IndexParams: []*commonpb.KeyValuePair{}},
|
|
},
|
|
}
|
|
collectionProps := []*commonpb.KeyValuePair{
|
|
{Key: common.WarmupScalarIndexKey, Value: common.WarmupDisable},
|
|
}
|
|
|
|
applyIndexWarmupSetting(loadInfo, schema, collectionProps)
|
|
|
|
indexInfo := loadInfo.GetIndexInfos()[0]
|
|
warmup, exist := common.GetWarmupPolicy(indexInfo.IndexParams...)
|
|
assert.True(t, exist)
|
|
assert.Equal(t, common.WarmupDisable, warmup, "collection-level setting should override autoWarmupForNonPKIsolationCollection")
|
|
})
|
|
}
|
|
|
|
func TestUtils(t *testing.T) {
|
|
suite.Run(t, new(UtilsSuite))
|
|
}
|