enhance: Check load fields for previous loaded collection (#35905)

Related to #35415 This PR make querycoord report error when load request tries to update load fields list, which is currently not supported. Signed-off-by: Congqi Xia <[email protected]>
milvus-io · Sep 2, 2024 · 3698c53 · 3698c53
1 parent 4641fd9
commit 3698c53
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 0 deletions.
diff --git a/internal/querycoordv2/job/job_load.go b/internal/querycoordv2/job/job_load.go
@@ -19,6 +19,7 @@ package job
 import (
 	"context"
 	"fmt"
+	"reflect"
 	"time"
 
 	"github.com/cockroachdb/errors"
@@ -104,6 +105,14 @@ func (job *LoadCollectionJob) PreExecute() error {
 		return merr.WrapErrParameterInvalid(collection.GetReplicaNumber(), req.GetReplicaNumber(), "can't change the replica number for loaded collection")
 	}
 
+	if !reflect.DeepEqual(collection.GetLoadFields(), req.GetLoadFields()) {
+		log.Warn("collection with different load field list exists, release this collection first before chaning its replica number",
+			zap.Int64s("loadedFieldIDs", collection.GetLoadFields()),
+			zap.Int64s("reqFieldIDs", req.GetLoadFields()),
+		)
+		return merr.WrapErrParameterInvalid(collection.GetLoadFields(), req.GetLoadFields(), "can't change the load field list for loaded collection")
+	}
+
 	return nil
 }
 
@@ -289,6 +298,14 @@ func (job *LoadPartitionJob) PreExecute() error {
 		return merr.WrapErrParameterInvalid(collection.GetReplicaNumber(), req.GetReplicaNumber(), "can't change the replica number for loaded partitions")
 	}
 
+	if !reflect.DeepEqual(collection.GetLoadFields(), req.GetLoadFields()) {
+		log.Warn("collection with different load field list exists, release this collection first before chaning its replica number",
+			zap.Int64s("loadedFieldIDs", collection.GetLoadFields()),
+			zap.Int64s("reqFieldIDs", req.GetLoadFields()),
+		)
+		return merr.WrapErrParameterInvalid(collection.GetLoadFields(), req.GetLoadFields(), "can't change the load field list for loaded collection")
+	}
+
 	return nil
 }
 

diff --git a/internal/querycoordv2/job/job_test.go b/internal/querycoordv2/job/job_test.go
@@ -307,6 +307,32 @@ func (suite *JobSuite) TestLoadCollection() {
 		suite.ErrorIs(err, merr.ErrParameterInvalid)
 	}
 
+	// Test load existed collection with different load fields
+	for _, collection := range suite.collections {
+		if suite.loadTypes[collection] != querypb.LoadType_LoadCollection {
+			continue
+		}
+		req := &querypb.LoadCollectionRequest{
+			CollectionID: collection,
+			LoadFields:   []int64{100, 101},
+		}
+		job := NewLoadCollectionJob(
+			ctx,
+			req,
+			suite.dist,
+			suite.meta,
+			suite.broker,
+			suite.cluster,
+			suite.targetMgr,
+			suite.targetObserver,
+			suite.collectionObserver,
+			suite.nodeMgr,
+		)
+		suite.scheduler.Add(job)
+		err := job.Wait()
+		suite.ErrorIs(err, merr.ErrParameterInvalid)
+	}
+
 	// Test load partition while collection exists
 	for _, collection := range suite.collections {
 		if suite.loadTypes[collection] != querypb.LoadType_LoadCollection {
@@ -514,6 +540,34 @@ func (suite *JobSuite) TestLoadPartition() {
 		suite.ErrorIs(err, merr.ErrParameterInvalid)
 	}
 
+	// Test load partition with different load fields
+	for _, collection := range suite.collections {
+		if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {
+			continue
+		}
+
+		req := &querypb.LoadPartitionsRequest{
+			CollectionID: collection,
+			PartitionIDs: suite.partitions[collection],
+			LoadFields:   []int64{100, 101},
+		}
+		job := NewLoadPartitionJob(
+			ctx,
+			req,
+			suite.dist,
+			suite.meta,
+			suite.broker,
+			suite.cluster,
+			suite.targetMgr,
+			suite.targetObserver,
+			suite.collectionObserver,
+			suite.nodeMgr,
+		)
+		suite.scheduler.Add(job)
+		err := job.Wait()
+		suite.ErrorIs(err, merr.ErrParameterInvalid)
+	}
+
 	// Test load partition with more partition
 	for _, collection := range suite.collections {
 		if suite.loadTypes[collection] != querypb.LoadType_LoadPartition {