admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-07 01:28:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: disallow expr when partition key isolation is enabled (#35031)

Browse Source 
issue: #34336 

Signed-off-by: Patrick Weizhi Xu <weizhi.xu@zilliz.com>
(cherry picked from commit 2889481ce9a14230e9c7f1c8f9c3c1decde77e03)
This commit is contained in:
Patrick Weizhi Xu 2024-07-29 14:21:50 +08:00 committed by GitHub
parent 0e41f104c5
commit 3c7f73137e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 51 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
internal/proxy/task_search.go
View File

@ -367,6 +367,11 @@ func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error {
// set PartitionIDs for sub search // set PartitionIDs for sub search
if t.partitionKeyMode { if t.partitionKeyMode {
// isolatioin has tighter constraint, check first
mvErr := setQueryInfoIfMvEnable(queryInfo, t, plan)
if mvErr != nil {
return mvErr
}
partitionIDs, err2 := t.tryParsePartitionIDsFromPlan(plan) partitionIDs, err2 := t.tryParsePartitionIDsFromPlan(plan)
if err2 != nil { if err2 != nil {
return err2 return err2
@ -374,10 +379,6 @@ func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error {
if len(partitionIDs) > 0 { if len(partitionIDs) > 0 {
internalSubReq.PartitionIDs = partitionIDs internalSubReq.PartitionIDs = partitionIDs
t.partitionIDsSet.Upsert(partitionIDs...) t.partitionIDsSet.Upsert(partitionIDs...)
mvErr := setQueryInfoIfMvEnable(queryInfo, t, plan)
if mvErr != nil {
return mvErr
}
} }
} else { } else {
internalSubReq.PartitionIDs = t.SearchRequest.GetPartitionIDs() internalSubReq.PartitionIDs = t.SearchRequest.GetPartitionIDs()
@ -427,16 +428,17 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error {
t.SearchRequest.Offset = offset t.SearchRequest.Offset = offset
if t.partitionKeyMode { if t.partitionKeyMode {
// isolatioin has tighter constraint, check first
mvErr := setQueryInfoIfMvEnable(queryInfo, t, plan)
if mvErr != nil {
return mvErr
}
partitionIDs, err2 := t.tryParsePartitionIDsFromPlan(plan) partitionIDs, err2 := t.tryParsePartitionIDsFromPlan(plan)
if err2 != nil { if err2 != nil {
return err2 return err2
} }
if len(partitionIDs) > 0 { if len(partitionIDs) > 0 {
t.SearchRequest.PartitionIDs = partitionIDs t.SearchRequest.PartitionIDs = partitionIDs
mvErr := setQueryInfoIfMvEnable(queryInfo, t, plan)
if mvErr != nil {
return mvErr
}
} }
} }

24
internal/proxy/task_search_test.go
View File

@ -2756,9 +2756,12 @@ func (s *MaterializedViewTestSuite) TestMvEnabledPartitionKeyOnVarChar() {
} }
func (s *MaterializedViewTestSuite) TestMvEnabledPartitionKeyOnVarCharWithIsolation() { func (s *MaterializedViewTestSuite) TestMvEnabledPartitionKeyOnVarCharWithIsolation() {
isAdanceds := []bool{true, false}
for _, isAdvanced := range isAdanceds {
task := s.getSearchTask() task := s.getSearchTask()
task.enableMaterializedView = true task.enableMaterializedView = true
task.request.Dsl = testVarCharField + " == \"a\"" task.request.Dsl = testVarCharField + " == \"a\""
task.IsAdvanced = isAdvanced
schema := ConstructCollectionSchemaWithPartitionKey(s.colName, s.fieldName2Types, testInt64Field, testVarCharField, false) schema := ConstructCollectionSchemaWithPartitionKey(s.colName, s.fieldName2Types, testInt64Field, testVarCharField, false)
schemaInfo := newSchemaInfo(schema) schemaInfo := newSchemaInfo(schema)
s.mockMetaCache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(schemaInfo, nil) s.mockMetaCache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(schemaInfo, nil)
@ -2769,18 +2772,35 @@ func (s *MaterializedViewTestSuite) TestMvEnabledPartitionKeyOnVarCharWithIsolat
s.NotZero(len(task.queryInfos)) s.NotZero(len(task.queryInfos))
s.Equal(true, task.queryInfos[0].MaterializedViewInvolved) s.Equal(true, task.queryInfos[0].MaterializedViewInvolved)
} }
}
func (s *MaterializedViewTestSuite) TestMvEnabledPartitionKeyOnVarCharWithIsolationInvalid() { func (s *MaterializedViewTestSuite) TestMvEnabledPartitionKeyOnVarCharWithIsolationInvalid() {
isAdanceds := []bool{true, false}
for _, isAdvanced := range isAdanceds {
task := s.getSearchTask() task := s.getSearchTask()
task.enableMaterializedView = true task.enableMaterializedView = true
task.IsAdvanced = isAdvanced
task.request.Dsl = testVarCharField + " in [\"a\", \"b\"]" task.request.Dsl = testVarCharField + " in [\"a\", \"b\"]"
schema := ConstructCollectionSchemaWithPartitionKey(s.colName, s.fieldName2Types, testInt64Field, testVarCharField, false) schema := ConstructCollectionSchemaWithPartitionKey(s.colName, s.fieldName2Types, testInt64Field, testVarCharField, false)
schemaInfo := newSchemaInfo(schema) schemaInfo := newSchemaInfo(schema)
s.mockMetaCache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(schemaInfo, nil) s.mockMetaCache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(schemaInfo, nil)
s.mockMetaCache.EXPECT().GetPartitionsIndex(mock.Anything, mock.Anything, mock.Anything).Return([]string{"partition_1", "partition_2"}, nil)
s.mockMetaCache.EXPECT().GetPartitions(mock.Anything, mock.Anything, mock.Anything).Return(map[string]int64{"partition_1": 1, "partition_2": 2}, nil)
s.ErrorContains(task.PreExecute(s.ctx), "partition key isolation does not support IN") s.ErrorContains(task.PreExecute(s.ctx), "partition key isolation does not support IN")
} }
}
func (s *MaterializedViewTestSuite) TestMvEnabledPartitionKeyOnVarCharWithIsolationInvalidOr() {
isAdanceds := []bool{true, false}
for _, isAdvanced := range isAdanceds {
task := s.getSearchTask()
task.enableMaterializedView = true
task.IsAdvanced = isAdvanced
task.request.Dsl = testVarCharField + " == \"a\" || " + testVarCharField + " == \"b\""
schema := ConstructCollectionSchemaWithPartitionKey(s.colName, s.fieldName2Types, testInt64Field, testVarCharField, false)
schemaInfo := newSchemaInfo(schema)
s.mockMetaCache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(schemaInfo, nil)
s.ErrorContains(task.PreExecute(s.ctx), "partition key isolation does not support OR")
}
}
func TestMaterializedView(t *testing.T) { func TestMaterializedView(t *testing.T) {
suite.Run(t, new(MaterializedViewTestSuite)) suite.Run(t, new(MaterializedViewTestSuite))