test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-11-30 02:48:45 +08:00
Code Issues Actions 11 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
if t.partitionKeyMode {
// isolatioin has tighter constraint, check first
mvErr := setQueryInfoIfMvEnable(queryInfo, t, plan)
if mvErr != nil {
return mvErr
}
partitionIDs, err2 := t.tryParsePartitionIDsFromPlan(plan)
if err2 != nil {
return err2
@ -374,10 +379,6 @@ func (t *searchTask) initAdvancedSearchRequest(ctx context.Context) error {
if len(partitionIDs) > 0 {
internalSubReq.PartitionIDs = partitionIDs
t.partitionIDsSet.Upsert(partitionIDs...)
mvErr := setQueryInfoIfMvEnable(queryInfo, t, plan)
if mvErr != nil {
return mvErr
}
}
} else {
internalSubReq.PartitionIDs = t.SearchRequest.GetPartitionIDs()
@ -427,16 +428,17 @@ func (t *searchTask) initSearchRequest(ctx context.Context) error {
t.SearchRequest.Offset = offset
if t.partitionKeyMode {
// isolatioin has tighter constraint, check first
mvErr := setQueryInfoIfMvEnable(queryInfo, t, plan)
if mvErr != nil {
return mvErr
}
partitionIDs, err2 := t.tryParsePartitionIDsFromPlan(plan)
if err2 != nil {
return err2
}
if len(partitionIDs) > 0 {
t.SearchRequest.PartitionIDs = partitionIDs
mvErr := setQueryInfoIfMvEnable(queryInfo, t, plan)
if mvErr != nil {
return mvErr
}
}
}

62
internal/proxy/task_search_test.go
View File

@ -2756,30 +2756,50 @@ func (s *MaterializedViewTestSuite) TestMvEnabledPartitionKeyOnVarChar() {
}
func (s *MaterializedViewTestSuite) TestMvEnabledPartitionKeyOnVarCharWithIsolation() {
task := s.getSearchTask()
task.enableMaterializedView = true
task.request.Dsl = testVarCharField + " == \"a\""
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.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)
err := task.PreExecute(s.ctx)
s.NoError(err)
s.NotZero(len(task.queryInfos))
s.Equal(true, task.queryInfos[0].MaterializedViewInvolved)
isAdanceds := []bool{true, false}
for _, isAdvanced := range isAdanceds {
task := s.getSearchTask()
task.enableMaterializedView = true
task.request.Dsl = testVarCharField + " == \"a\""
task.IsAdvanced = isAdvanced
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.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)
err := task.PreExecute(s.ctx)
s.NoError(err)
s.NotZero(len(task.queryInfos))
s.Equal(true, task.queryInfos[0].MaterializedViewInvolved)
}
}
func (s *MaterializedViewTestSuite) TestMvEnabledPartitionKeyOnVarCharWithIsolationInvalid() {
task := s.getSearchTask()
task.enableMaterializedView = true
task.request.Dsl = testVarCharField + " in [\"a\", \"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.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")
isAdanceds := []bool{true, false}
for _, isAdvanced := range isAdanceds {
task := s.getSearchTask()
task.enableMaterializedView = true
task.IsAdvanced = isAdvanced
task.request.Dsl = testVarCharField + " in [\"a\", \"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 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) {