diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index e1c0c1958a..5f40f6a03b 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -289,6 +289,10 @@ cd ../../../ make unittest ``` +The Apple Silicon compose file uses the Milvus Harbor Docker Hub proxy for +Pulsar by default. To pull Pulsar directly from Docker Hub instead, set +`PULSAR_IMAGE=milvusdb/pulsar:v2.8.2-m1`. + For others: ```shell diff --git a/client/milvusclient/iterator_option.go b/client/milvusclient/iterator_option.go index e44a1429ab..ff632fd988 100644 --- a/client/milvusclient/iterator_option.go +++ b/client/milvusclient/iterator_option.go @@ -70,6 +70,11 @@ func (opt *searchIteratorOption) WithPartitions(partitionNames ...string) *searc return opt } +func (opt *searchIteratorOption) WithNamespace(namespace string) *searchIteratorOption { + opt.searchOption.WithNamespace(namespace) + return opt +} + func (opt *searchIteratorOption) WithFilter(expr string) *searchIteratorOption { opt.annRequest.WithFilter(expr) return opt @@ -166,6 +171,7 @@ type QueryIteratorOption interface { type queryIteratorOption struct { collectionName string partitionNames []string + namespace *string outputFields []string expr string batchSize int @@ -178,6 +184,7 @@ func (opt *queryIteratorOption) Request() (*milvuspb.QueryRequest, error) { return &milvuspb.QueryRequest{ CollectionName: opt.collectionName, PartitionNames: opt.partitionNames, + Namespace: opt.namespace, OutputFields: opt.outputFields, Expr: opt.expr, ConsistencyLevel: opt.consistencyLevel.CommonConsistencyLevel(), @@ -211,6 +218,11 @@ func (opt *queryIteratorOption) WithPartitions(partitionNames ...string) *queryI return opt } +func (opt *queryIteratorOption) WithNamespace(namespace string) *queryIteratorOption { + opt.namespace = &namespace + return opt +} + func (opt *queryIteratorOption) WithFilter(expr string) *queryIteratorOption { opt.expr = expr return opt diff --git a/client/milvusclient/iterator_test.go b/client/milvusclient/iterator_test.go index 656b37c8b8..11e2301f5a 100644 --- a/client/milvusclient/iterator_test.go +++ b/client/milvusclient/iterator_test.go @@ -47,6 +47,18 @@ func (s *SearchIteratorSuite) SetupSuite() { WithField(entity.NewField().WithName("Vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128)) } +func (s *SearchIteratorSuite) TestSearchIteratorOptionWithNamespace() { + namespace := "tenant_a" + + opt := NewSearchIteratorOption("coll", entity.FloatVector(lo.RepeatBy(128, func(_ int) float32 { + return rand.Float32() + }))).WithNamespace(namespace) + req, err := opt.SearchOption().Request() + + s.Require().NoError(err) + s.Equal(namespace, req.GetNamespace()) +} + func (s *SearchIteratorSuite) TestSearchIteratorInit() { ctx := context.Background() s.Run("success", func() { @@ -431,6 +443,15 @@ func (s *QueryIteratorSuite) SetupSuite() { WithField(entity.NewField().WithName("Name").WithDataType(entity.FieldTypeVarChar).WithMaxLength(256)) } +func (s *QueryIteratorSuite) TestQueryIteratorOptionWithNamespace() { + namespace := "tenant_a" + + req, err := NewQueryIteratorOption("coll").WithNamespace(namespace).Request() + + s.Require().NoError(err) + s.Equal(namespace, req.GetNamespace()) +} + func (s *QueryIteratorSuite) TestQueryIteratorInit() { ctx := context.Background() s.Run("success", func() { diff --git a/client/milvusclient/read_option_test.go b/client/milvusclient/read_option_test.go index 9fd948a634..451f394fc7 100644 --- a/client/milvusclient/read_option_test.go +++ b/client/milvusclient/read_option_test.go @@ -90,6 +90,29 @@ func (s *SearchOptionSuite) TestBasic() { s.Error(err) } +func (s *SearchOptionSuite) TestWithNamespace() { + collName := "namespace_read_option" + namespace := "tenant_a" + + searchReq, err := NewSearchOption(collName, 10, []entity.Vector{entity.FloatVector([]float32{0.1, 0.2})}). + WithNamespace(namespace). + Request() + s.Require().NoError(err) + s.Equal(namespace, searchReq.GetNamespace()) + + queryReq, err := NewQueryOption(collName). + WithNamespace(namespace). + Request() + s.Require().NoError(err) + s.Equal(namespace, queryReq.GetNamespace()) + + hybridReq, err := NewHybridSearchOption(collName, 10, NewAnnRequest("vector", 10, entity.FloatVector([]float32{0.1, 0.2}))). + WithNamespace(namespace). + HybridRequest() + s.Require().NoError(err) + s.Equal(namespace, hybridReq.GetNamespace()) +} + func (s *SearchOptionSuite) TestPlaceHolder() { type testCase struct { tag string diff --git a/client/milvusclient/read_options.go b/client/milvusclient/read_options.go index f0db347caa..96b87c7361 100644 --- a/client/milvusclient/read_options.go +++ b/client/milvusclient/read_options.go @@ -59,6 +59,7 @@ type searchOption struct { annRequest *AnnRequest collectionName string partitionNames []string + namespace *string outputFields []string searchAggregation *SearchAggregation consistencyLevel entity.ConsistencyLevel @@ -333,6 +334,7 @@ func (opt *searchOption) Request() (*milvuspb.SearchRequest, error) { request.CollectionName = opt.collectionName request.PartitionNames = opt.partitionNames + request.Namespace = opt.namespace request.ConsistencyLevel = commonpb.ConsistencyLevel(opt.consistencyLevel) request.UseDefaultConsistency = opt.useDefaultConsistencyLevel request.OutputFields = opt.outputFields @@ -367,6 +369,11 @@ func (opt *searchOption) WithPartitions(partitionNames ...string) *searchOption return opt } +func (opt *searchOption) WithNamespace(namespace string) *searchOption { + opt.namespace = &namespace + return opt +} + func (opt *searchOption) WithFilter(expr string) *searchOption { opt.annRequest.WithFilter(expr) return opt @@ -525,6 +532,7 @@ type HybridSearchOption interface { type hybridSearchOption struct { collectionName string partitionNames []string + namespace *string reqs []*AnnRequest @@ -554,6 +562,11 @@ func (opt *hybridSearchOption) WithPartitions(partitions ...string) *hybridSearc return opt } +func (opt *hybridSearchOption) WithNamespace(namespace string) *hybridSearchOption { + opt.namespace = &namespace + return opt +} + func (opt *hybridSearchOption) WithOutputFields(outputFields ...string) *hybridSearchOption { opt.outputFields = outputFields return opt @@ -596,6 +609,7 @@ func (opt *hybridSearchOption) HybridRequest() (*milvuspb.HybridSearchRequest, e r := &milvuspb.HybridSearchRequest{ CollectionName: opt.collectionName, PartitionNames: opt.partitionNames, + Namespace: opt.namespace, Requests: requests, UseDefaultConsistency: opt.useDefaultConsistency, ConsistencyLevel: commonpb.ConsistencyLevel(opt.consistencyLevel), @@ -629,6 +643,7 @@ type QueryOption interface { type queryOption struct { collectionName string partitionNames []string + namespace *string queryParams map[string]string outputFields []string consistencyLevel entity.ConsistencyLevel @@ -641,6 +656,7 @@ func (opt *queryOption) Request() (*milvuspb.QueryRequest, error) { req := &milvuspb.QueryRequest{ CollectionName: opt.collectionName, PartitionNames: opt.partitionNames, + Namespace: opt.namespace, OutputFields: opt.outputFields, Expr: opt.expr, @@ -703,6 +719,11 @@ func (opt *queryOption) WithPartitions(partitionNames ...string) *queryOption { return opt } +func (opt *queryOption) WithNamespace(namespace string) *queryOption { + opt.namespace = &namespace + return opt +} + func (opt *queryOption) WithIDs(ids column.Column) *queryOption { opt.expr = pks2Expr(ids) return opt diff --git a/client/milvusclient/write_option_test.go b/client/milvusclient/write_option_test.go index 6db7a33b74..a23b9a30e5 100644 --- a/client/milvusclient/write_option_test.go +++ b/client/milvusclient/write_option_test.go @@ -195,6 +195,59 @@ func (s *ColumnBasedDataOptionSuite) TestNewStructSubColumnErrors() { } } +func (s *ColumnBasedDataOptionSuite) TestWithNamespace() { + collName := "namespace_write_option" + namespace := "tenant_a" + coll := &entity.Collection{ + Schema: entity.NewSchema().WithName(collName). + WithField(entity.NewField().WithName("id").WithDataType(entity.FieldTypeInt64)), + } + + insertOpt := NewColumnBasedInsertOption(collName, column.NewColumnInt64("id", []int64{1})). + WithNamespace(namespace) + insertReq, err := insertOpt.InsertRequest(coll) + s.Require().NoError(err) + s.Equal(namespace, insertReq.GetNamespace()) + + upsertOpt := NewColumnBasedInsertOption(collName, column.NewColumnInt64("id", []int64{1})). + WithNamespace(namespace) + upsertReq, err := upsertOpt.UpsertRequest(coll) + s.Require().NoError(err) + s.Equal(namespace, upsertReq.GetNamespace()) +} + +func (s *ColumnBasedDataOptionSuite) TestRowBasedWithNamespaceKeepsRows() { + collName := "namespace_row_write_option" + namespace := "tenant_a" + partition := "partition_a" + coll := &entity.Collection{ + Schema: entity.NewSchema().WithName(collName). + WithField(entity.NewField().WithName("id").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)). + WithField(entity.NewField().WithName("name").WithDataType(entity.FieldTypeVarChar).WithMaxLength(64)), + } + rows := []any{map[string]any{"id": int64(1), "name": "alice"}} + + var insertOpt InsertOption = NewRowBasedInsertOption(collName, rows...). + WithPartition(partition). + WithNamespace(namespace) + insertReq, err := insertOpt.InsertRequest(coll) + s.Require().NoError(err) + s.Equal(partition, insertReq.GetPartitionName()) + s.Equal(namespace, insertReq.GetNamespace()) + s.EqualValues(1, insertReq.GetNumRows()) + s.Len(insertReq.GetFieldsData(), 2) + + var upsertOpt UpsertOption = NewRowBasedInsertOption(collName, rows...). + WithPartition(partition). + WithNamespace(namespace) + upsertReq, err := upsertOpt.UpsertRequest(coll) + s.Require().NoError(err) + s.Equal(partition, upsertReq.GetPartitionName()) + s.Equal(namespace, upsertReq.GetNamespace()) + s.EqualValues(1, upsertReq.GetNumRows()) + s.Len(upsertReq.GetFieldsData(), 2) +} + func TestRowBasedDataOption(t *testing.T) { suite.Run(t, new(ColumnBasedDataOptionSuite)) } @@ -210,6 +263,14 @@ func (s *DeleteOptionSuite) TestBasic() { s.Equal(collectionName, opt.Request().GetCollectionName()) } +func (s *DeleteOptionSuite) TestWithNamespace() { + collectionName := fmt.Sprintf("coll_%s", s.randString(6)) + namespace := "tenant_a" + + req := NewDeleteOption(collectionName).WithNamespace(namespace).Request() + s.Equal(namespace, req.GetNamespace()) +} + func TestDeleteOption(t *testing.T) { suite.Run(t, new(DeleteOptionSuite)) } diff --git a/client/milvusclient/write_options.go b/client/milvusclient/write_options.go index a7f4dda37c..f45b1af61b 100644 --- a/client/milvusclient/write_options.go +++ b/client/milvusclient/write_options.go @@ -51,6 +51,7 @@ var ( type columnBasedDataOption struct { collName string partitionName string + namespace *string columns []column.Column partialUpdate bool @@ -379,6 +380,14 @@ func (opt *columnBasedDataOption) WithPartition(partitionName string) *columnBas return opt } +// WithNamespace scopes the write to a collection namespace. Primary keys are +// still collection-scoped for delete/upsert tombstones, so callers must keep +// primary keys unique across namespaces in the same collection. +func (opt *columnBasedDataOption) WithNamespace(namespace string) *columnBasedDataOption { + opt.namespace = &namespace + return opt +} + func (opt *columnBasedDataOption) WithPartialUpdate(partialUpdate bool) *columnBasedDataOption { opt.partialUpdate = partialUpdate return opt @@ -452,6 +461,7 @@ func (opt *columnBasedDataOption) InsertRequest(coll *entity.Collection) (*milvu return &milvuspb.InsertRequest{ CollectionName: opt.collName, PartitionName: opt.partitionName, + Namespace: opt.namespace, FieldsData: fieldsData, NumRows: uint32(rowNum), SchemaTimestamp: coll.UpdateTimestamp, @@ -477,6 +487,7 @@ func (opt *columnBasedDataOption) UpsertRequest(coll *entity.Collection) (*milvu return &milvuspb.UpsertRequest{ CollectionName: opt.collName, PartitionName: opt.partitionName, + Namespace: opt.namespace, FieldsData: fieldsData, NumRows: uint32(rowNum), SchemaTimestamp: coll.UpdateTimestamp, @@ -509,6 +520,36 @@ func NewRowBasedInsertOption(collName string, rows ...any) *rowBasedDataOption { } } +func (opt *rowBasedDataOption) WithPartition(partitionName string) *rowBasedDataOption { + opt.columnBasedDataOption.WithPartition(partitionName) + return opt +} + +func (opt *rowBasedDataOption) WithNamespace(namespace string) *rowBasedDataOption { + opt.columnBasedDataOption.WithNamespace(namespace) + return opt +} + +func (opt *rowBasedDataOption) WithPartialUpdate(partialUpdate bool) *rowBasedDataOption { + opt.columnBasedDataOption.WithPartialUpdate(partialUpdate) + return opt +} + +func (opt *rowBasedDataOption) WithArrayAppend(fieldName string) *rowBasedDataOption { + opt.columnBasedDataOption.WithArrayAppend(fieldName) + return opt +} + +func (opt *rowBasedDataOption) WithArrayRemove(fieldName string) *rowBasedDataOption { + opt.columnBasedDataOption.WithArrayRemove(fieldName) + return opt +} + +func (opt *rowBasedDataOption) WithFieldPartialOp(fieldName string, op schemapb.FieldPartialUpdateOp_OpType) *rowBasedDataOption { + opt.columnBasedDataOption.WithFieldPartialOp(fieldName, op) + return opt +} + func (opt *rowBasedDataOption) InsertRequest(coll *entity.Collection) (*milvuspb.InsertRequest, error) { columns, err := row.AnyToColumns(opt.rows, opt.keepAutoIDPk, coll.Schema) if err != nil { @@ -522,6 +563,7 @@ func (opt *rowBasedDataOption) InsertRequest(coll *entity.Collection) (*milvuspb return &milvuspb.InsertRequest{ CollectionName: opt.collName, PartitionName: opt.partitionName, + Namespace: opt.namespace, FieldsData: fieldsData, NumRows: uint32(rowNum), }, nil @@ -545,6 +587,7 @@ func (opt *rowBasedDataOption) UpsertRequest(coll *entity.Collection) (*milvuspb return &milvuspb.UpsertRequest{ CollectionName: opt.collName, PartitionName: opt.partitionName, + Namespace: opt.namespace, FieldsData: fieldsData, NumRows: uint32(rowNum), PartialUpdate: partialUpdate, @@ -586,6 +629,7 @@ type DeleteOption interface { type deleteOption struct { collectionName string partitionName string + namespace *string expr string } @@ -593,6 +637,7 @@ func (opt *deleteOption) Request() *milvuspb.DeleteRequest { return &milvuspb.DeleteRequest{ CollectionName: opt.collectionName, PartitionName: opt.partitionName, + Namespace: opt.namespace, Expr: opt.expr, } } @@ -617,6 +662,14 @@ func (opt *deleteOption) WithPartition(partitionName string) *deleteOption { return opt } +// WithNamespace scopes the delete request to a collection namespace. Delete +// tombstones are primary-key based, so callers must keep primary keys unique +// across namespaces in the same collection. +func (opt *deleteOption) WithNamespace(namespace string) *deleteOption { + opt.namespace = &namespace + return opt +} + func NewDeleteOption(collectionName string) *deleteOption { return &deleteOption{collectionName: collectionName} } diff --git a/client/milvusclient/write_options_partial_op_test.go b/client/milvusclient/write_options_partial_op_test.go index 2c6d5a299e..fc91b12511 100644 --- a/client/milvusclient/write_options_partial_op_test.go +++ b/client/milvusclient/write_options_partial_op_test.go @@ -153,11 +153,12 @@ func TestRowBasedUpsertEmitsFieldOps(t *testing.T) { map[string]any{"id": int64(1), "tags": []int64{10}}, map[string]any{"id": int64(2), "tags": []int64{20, 30}}, } - opt := NewRowBasedInsertOption(coll.Name, rows...) - opt.WithArrayAppend("tags") + opt := NewRowBasedInsertOption(coll.Name, rows...). + WithArrayAppend("tags") req, err := opt.UpsertRequest(coll) require.NoError(t, err) + assert.EqualValues(t, len(rows), req.GetNumRows()) assert.True(t, req.GetPartialUpdate()) tagsOp := findOp(req.GetFieldOps(), "tags") require.NotNil(t, tagsOp) diff --git a/deployments/docker/dev/docker-compose-apple-silicon.yml b/deployments/docker/dev/docker-compose-apple-silicon.yml index f0a0f26ff0..1505a7ee25 100644 --- a/deployments/docker/dev/docker-compose-apple-silicon.yml +++ b/deployments/docker/dev/docker-compose-apple-silicon.yml @@ -1,5 +1,3 @@ -version: '3.5' - services: etcd: image: quay.io/coreos/etcd:v3.5.25 @@ -17,7 +15,7 @@ services: - "4001:4001" pulsar: - image: milvusdb/pulsar:v2.8.2-m1 + image: ${PULSAR_IMAGE:-harbor.milvus.io/dockerhub/milvusdb/pulsar:v2.8.2-m1} volumes: - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/pulsar:/pulsar/data environment: diff --git a/internal/proxy/impl.go b/internal/proxy/impl.go index 283718e406..5a937fceb5 100644 --- a/internal/proxy/impl.go +++ b/internal/proxy/impl.go @@ -3027,6 +3027,7 @@ func (node *Proxy) search(ctx context.Context, request *milvuspb.SearchRequest, shardClientMgr: node.shardMgr, enableMaterializedView: node.enableMaterializedView, mustUsePartitionKey: Params.ProxyCfg.MustUsePartitionKey.GetAsBool(), + chMgr: node.chMgr, } succeeded := false @@ -3254,6 +3255,7 @@ func (node *Proxy) hybridSearch(ctx context.Context, request *milvuspb.HybridSea lb: node.lbPolicy, shardClientMgr: node.shardMgr, mustUsePartitionKey: Params.ProxyCfg.MustUsePartitionKey.GetAsBool(), + chMgr: node.chMgr, } succeeded := false @@ -3529,6 +3531,7 @@ func (node *Proxy) handleIfSearchByPK(ctx context.Context, request *milvuspb.Sea GuaranteeTimestamp: request.GuaranteeTimestamp, ConsistencyLevel: request.ConsistencyLevel, UseDefaultConsistency: request.UseDefaultConsistency, + Namespace: request.Namespace, } // Create queryTask to execute the retrieval @@ -3552,6 +3555,7 @@ func (node *Proxy) handleIfSearchByPK(ctx context.Context, request *milvuspb.Sea mustUsePartitionKey: Params.ProxyCfg.MustUsePartitionKey.GetAsBool(), // reQuery defaults to false - we need full query processing: // partition conversion, struct field reconstruction, timestamp handling etc + chMgr: node.chMgr, } // Execute query @@ -3828,6 +3832,7 @@ func (node *Proxy) Query(ctx context.Context, request *milvuspb.QueryRequest) (* lb: node.lbPolicy, shardclientMgr: node.shardMgr, mustUsePartitionKey: Params.ProxyCfg.MustUsePartitionKey.GetAsBool(), + chMgr: node.chMgr, } subLabel := GetCollectionRateSubLabel(request) diff --git a/internal/proxy/impl_test.go b/internal/proxy/impl_test.go index 371aa60d84..fc9c6887ed 100644 --- a/internal/proxy/impl_test.go +++ b/internal/proxy/impl_test.go @@ -31,6 +31,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/trace" "google.golang.org/grpc" "google.golang.org/grpc/metadata" "google.golang.org/protobuf/proto" @@ -46,6 +47,7 @@ import ( "github.com/milvus-io/milvus/internal/mocks/distributed/mock_streaming" "github.com/milvus-io/milvus/internal/proxy/shardclient" "github.com/milvus-io/milvus/internal/util/dependency" + "github.com/milvus-io/milvus/internal/util/segcore" "github.com/milvus-io/milvus/internal/util/sessionutil" "github.com/milvus-io/milvus/pkg/v3/common" "github.com/milvus-io/milvus/pkg/v3/proto/datapb" @@ -2207,6 +2209,88 @@ func TestHandleIfSearchByPK_BM25Detection(t *testing.T) { }) } +func TestHandleIfSearchByPK_PreservesNamespaceInInternalQuery(t *testing.T) { + mockey.PatchConvey("TestHandleIfSearchByPK_PreservesNamespaceInInternalQuery", t, func() { + paramtable.Init() + originalCache := globalMetaCache + defer func() { globalMetaCache = originalCache }() + + namespace := "tenant_a" + schema := &schemapb.CollectionSchema{ + Name: "test_collection", + EnableNamespace: true, + Fields: []*schemapb.FieldSchema{ + {FieldID: 100, Name: "id", IsPrimaryKey: true, DataType: schemapb.DataType_Int64}, + { + FieldID: 101, + Name: "vec", + DataType: schemapb.DataType_FloatVector, + TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "2"}}, + }, + }, + } + + cache := NewMockCache(t) + cache.EXPECT(). + GetCollectionInfo(mock.Anything, "default", "test_collection", int64(0)). + Return(&collectionInfo{schema: newSchemaInfo(schema)}, nil) + globalMetaCache = cache + + var capturedNamespace *string + mockey.Mock((*Proxy).query).To(func(_ *Proxy, _ context.Context, qt *queryTask, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) { + capturedNamespace = qt.request.Namespace + return &milvuspb.QueryResults{ + Status: merr.Success(), + FieldsData: []*schemapb.FieldData{ + { + FieldName: "id", + FieldId: 100, + Type: schemapb.DataType_Int64, + Field: &schemapb.FieldData_Scalars{ + Scalars: &schemapb.ScalarField{ + Data: &schemapb.ScalarField_LongData{ + LongData: &schemapb.LongArray{Data: []int64{1}}, + }, + }, + }, + }, + { + FieldName: "vec", + FieldId: 101, + Type: schemapb.DataType_FloatVector, + Field: &schemapb.FieldData_Vectors{ + Vectors: &schemapb.VectorField{ + Dim: 2, + Data: &schemapb.VectorField_FloatVector{ + FloatVector: &schemapb.FloatArray{Data: []float32{0.1, 0.2}}, + }, + }, + }, + }, + }, + }, segcore.StorageCost{}, nil + }).Build() + + node := &Proxy{} + req := &milvuspb.SearchRequest{ + DbName: "default", + CollectionName: "test_collection", + Namespace: &namespace, + SearchInput: &milvuspb.SearchRequest_Ids{ + Ids: &schemapb.IDs{ + IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}}, + }, + }, + SearchParams: []*commonpb.KeyValuePair{{Key: AnnsFieldKey, Value: "vec"}}, + } + + _, err := node.handleIfSearchByPK(context.Background(), req) + assert.NoError(t, err) + require.NotNil(t, capturedNamespace) + assert.Equal(t, namespace, *capturedNamespace) + }) +} + func TestProxy_ManualCompaction_ExternalCollection(t *testing.T) { // Save and restore globalMetaCache cache := globalMetaCache diff --git a/internal/proxy/routing_table_hash_test.go b/internal/proxy/routing_table_hash_test.go index ff67ae5883..7167486cd7 100644 --- a/internal/proxy/routing_table_hash_test.go +++ b/internal/proxy/routing_table_hash_test.go @@ -84,6 +84,8 @@ func TestRepackDeleteMsgByHashPreservesModuloRouting(t *testing.T) { 2, "partition", "default", + nil, + nil, ) assert.NoError(t, err) @@ -125,6 +127,8 @@ func TestRepackDeleteMsgByHashReturnsRoutingErrorWithoutChannels(t *testing.T) { 2, "partition", "default", + nil, + nil, ) assert.ErrorIs(t, err, common.ErrRoutingTableNoValues) diff --git a/internal/proxy/search_pipeline.go b/internal/proxy/search_pipeline.go index 4724966c5e..ce075c50c2 100644 --- a/internal/proxy/search_pipeline.go +++ b/internal/proxy/search_pipeline.go @@ -1478,6 +1478,7 @@ func (op *requeryOperator) requery(ctx context.Context, span trace.Span, ids *sc preferredNodes: preferredNodes, fastSkip: true, reQuery: true, + chMgr: op.node.(*Proxy).chMgr, } queryResult, storageCost, err := op.node.(*Proxy).query(op.traceCtx, qt, span) if err != nil { diff --git a/internal/proxy/search_util.go b/internal/proxy/search_util.go index 8455f0ac21..6d8f268c08 100644 --- a/internal/proxy/search_util.go +++ b/internal/proxy/search_util.go @@ -1070,6 +1070,7 @@ func convertHybridSearchToSearch(req *milvuspb.HybridSearchRequest) *milvuspb.Se PartitionNames: req.GetPartitionNames(), OutputFields: req.GetOutputFields(), SearchParams: req.GetRankParams(), + Namespace: req.Namespace, TravelTimestamp: req.GetTravelTimestamp(), GuaranteeTimestamp: req.GetGuaranteeTimestamp(), Nq: 0, @@ -1079,7 +1080,6 @@ func convertHybridSearchToSearch(req *milvuspb.HybridSearchRequest) *milvuspb.Se SearchByPrimaryKeys: false, SubReqs: nil, FunctionScore: req.FunctionScore, - Namespace: req.Namespace, } for _, sub := range req.GetRequests() { diff --git a/internal/proxy/search_util_test.go b/internal/proxy/search_util_test.go new file mode 100644 index 0000000000..46be16d54e --- /dev/null +++ b/internal/proxy/search_util_test.go @@ -0,0 +1,39 @@ +// 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 ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb" +) + +func TestConvertHybridSearchToSearchKeepsNamespace(t *testing.T) { + namespace := "tenant_a" + + searchReq := convertHybridSearchToSearch(&milvuspb.HybridSearchRequest{ + CollectionName: "coll", + Namespace: &namespace, + Requests: []*milvuspb.SearchRequest{ + {Nq: 1}, + }, + }) + + assert.Equal(t, namespace, searchReq.GetNamespace()) +} diff --git a/internal/proxy/service_provider.go b/internal/proxy/service_provider.go index 2370b7f51a..b2f2a31c38 100644 --- a/internal/proxy/service_provider.go +++ b/internal/proxy/service_provider.go @@ -208,6 +208,7 @@ func (node *CachedProxyServiceProvider) DescribeCollection(ctx context.Context, }), StructArrayFields: cloneStructArrayFields(c.schema.StructArrayFields), EnableDynamicField: c.schema.EnableDynamicField, + EnableNamespace: c.schema.EnableNamespace, Properties: c.schema.Properties, Functions: c.schema.Functions, DbName: c.schema.DbName, diff --git a/internal/proxy/service_provider_test.go b/internal/proxy/service_provider_test.go index b593567699..7ebb56f758 100644 --- a/internal/proxy/service_provider_test.go +++ b/internal/proxy/service_provider_test.go @@ -91,6 +91,7 @@ func TestCachedProxyServiceProvider_DescribeCollection_FilterNamespaceField(t *t schema := &schemapb.CollectionSchema{ Name: collectionName, EnableDynamicField: true, + EnableNamespace: true, Fields: []*schemapb.FieldSchema{ { FieldID: common.StartOfUserFieldID, @@ -128,6 +129,7 @@ func TestCachedProxyServiceProvider_DescribeCollection_FilterNamespaceField(t *t }) assert.NoError(t, err) assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode()) + assert.True(t, resp.GetSchema().GetEnableNamespace()) fieldNames := make(map[string]struct{}) for _, f := range resp.GetSchema().GetFields() { diff --git a/internal/proxy/task_delete.go b/internal/proxy/task_delete.go index d15bed4c6f..1c527f255c 100644 --- a/internal/proxy/task_delete.go +++ b/internal/proxy/task_delete.go @@ -128,7 +128,14 @@ func (dt *deleteTask) getChannels() []pChan { } func (dt *deleteTask) PreExecute(ctx context.Context) error { - return nil + if dt.req.Namespace == nil { + return nil + } + schema, err := globalMetaCache.GetCollectionSchema(ctx, dt.req.GetDbName(), dt.req.GetCollectionName()) + if err != nil { + return err + } + return common.CheckNamespace(schema.CollectionSchema, dt.req.Namespace) } func (dt *deleteTask) PostExecute(ctx context.Context) error { @@ -151,12 +158,30 @@ func repackDeleteMsgByHash( partitionID int64, partitionName string, dbName string, + namespace *string, + schema *schemapb.CollectionSchema, ) (map[uint32][]*msgstream.DeleteMsg, int64, error) { maxSize := Params.PulsarCfg.MaxMessageSize.GetAsInt() - hashValues, err := typeutil.HashPK2Channels(primaryKeys, vChannels) + var hashValues []uint32 + // Delete tombstones are PK+timestamp based. Namespace can narrow routing, + // but it is not part of the tombstone identity; PKs must stay unique across + // namespaces in the same collection. + channelID, ok, err := namespaceShardingChannelID(schema, namespace, vChannels) if err != nil { return nil, 0, err } + if ok { + size := typeutil.GetSizeOfIDs(primaryKeys) + hashValues = make([]uint32, size) + for i := 0; i < size; i++ { + hashValues[i] = channelID + } + } else { + hashValues, err = typeutil.HashPK2Channels(primaryKeys, vChannels) + if err != nil { + return nil, 0, err + } + } // repack delete msg by dmChannel result := make(map[uint32][]*msgstream.DeleteMsg) lastMessageSize := map[uint32]int{} @@ -294,11 +319,11 @@ func (dr *deleteRunner) Init(ctx context.Context) error { if err := validateTextStorageV3Enabled(dr.schema.CollectionSchema); err != nil { return ErrWithLog(log, "TEXT field requires StorageV3", err) } - if namespacePartitionModeEnabled(dr.schema.CollectionSchema) { - partitionName, _, err := resolveNamespacePartitionName(dr.schema.CollectionSchema, dr.req.Namespace, dr.req.GetPartitionName()) - if err != nil { - return err - } + partitionName, namespaceAsPartition, err := resolveNamespacePartitionName(dr.schema.CollectionSchema, dr.req.Namespace, dr.req.GetPartitionName()) + if err != nil { + return err + } + if namespaceAsPartition { dr.req.PartitionName = partitionName } @@ -321,9 +346,22 @@ func (dr *deleteRunner) Init(ctx context.Context) error { return merr.WrapErrAsInputError(merr.WrapErrParameterInvalidMsg("delete plan can't be empty or always true : %s", dr.req.GetExpr())) } + dr.plan.Namespace = namespaceForPlan(dr.schema.CollectionSchema, dr.req.Namespace) // Set partitionIDs, could be empty if no partition name specified and no partition key partName := dr.req.GetPartitionName() - if dr.schema.IsPartitionKeyCollection() { + if namespacePartitionKeyMode(dr.schema.CollectionSchema) && dr.req.Namespace != nil { + if len(partName) > 0 { + return merr.WrapErrParameterInvalidMsg("not support manually specifying the partition names if namespace is used") + } + hashedPartitionNames, err := assignNamespacePartitionKey(ctx, dr.req.GetDbName(), dr.req.GetCollectionName(), dr.req.Namespace) + if err != nil { + return err + } + dr.partitionIDs, err = getPartitionIDs(ctx, dr.req.GetDbName(), dr.req.GetCollectionName(), hashedPartitionNames) + if err != nil { + return err + } + } else if dr.schema.IsPartitionKeyCollection() { if len(partName) > 0 { return merr.WrapErrParameterInvalidMsg("not support manually specifying the partition names if partition key mode is used") } @@ -561,13 +599,28 @@ func (dr *deleteRunner) complexDelete(ctx context.Context, plan *planpb.PlanNode return err } - err = dr.lb.Execute(ctx, shardclient.CollectionWorkLoad{ - Db: dr.req.GetDbName(), - CollectionName: dr.req.GetCollectionName(), - CollectionID: dr.collectionID, - Nq: 1, - Exec: dr.getStreamingQueryAndDelteFunc(plan), - }) + channelName, useNamespaceChannel, err := namespaceShardingChannel(dr.schema.CollectionSchema, dr.req.Namespace, dr.vChannels) + if err != nil { + return err + } + if useNamespaceChannel { + err = dr.lb.ExecuteWithRetry(ctx, shardclient.ChannelWorkload{ + Db: dr.req.GetDbName(), + CollectionName: dr.req.GetCollectionName(), + CollectionID: dr.collectionID, + Channel: channelName, + Nq: 1, + Exec: dr.getStreamingQueryAndDelteFunc(plan), + }) + } else { + err = dr.lb.Execute(ctx, shardclient.CollectionWorkLoad{ + Db: dr.req.GetDbName(), + CollectionName: dr.req.GetCollectionName(), + CollectionID: dr.collectionID, + Nq: 1, + Exec: dr.getStreamingQueryAndDelteFunc(plan), + }) + } dr.result.DeleteCnt = dr.count.Load() dr.result.Timestamp = dr.sessionTS.Load() if err != nil { diff --git a/internal/proxy/task_delete_streaming.go b/internal/proxy/task_delete_streaming.go index 5035d5e7d5..07f8db67a5 100644 --- a/internal/proxy/task_delete_streaming.go +++ b/internal/proxy/task_delete_streaming.go @@ -6,6 +6,7 @@ import ( "go.opentelemetry.io/otel" + "github.com/milvus-io/milvus-proto/go-api/v3/schemapb" "github.com/milvus-io/milvus/internal/distributed/streaming" "github.com/milvus-io/milvus/internal/util/hookutil" "github.com/milvus-io/milvus/pkg/v3/mlog" @@ -26,6 +27,17 @@ func (dt *deleteTask) Execute(ctx context.Context) (err error) { } dt.tr = timerecord.NewTimeRecorder(fmt.Sprintf("proxy execute delete %d", dt.ID())) + + var collectionSchema *schemapb.CollectionSchema + if dt.req.Namespace != nil || hookutil.IsClusterEncryptionEnabled() { + schema, err := globalMetaCache.GetCollectionSchema(ctx, dt.req.GetDbName(), dt.req.GetCollectionName()) + if err != nil { + mlog.Warn(ctx, "get collection schema from global meta cache failed", mlog.String("collectionName", dt.req.GetCollectionName()), mlog.Err(err)) + return err + } + collectionSchema = schema.CollectionSchema + } + result, numRows, err := repackDeleteMsgByHash( ctx, dt.primaryKeys, dt.vChannels, dt.idAllocator, @@ -33,6 +45,8 @@ func (dt *deleteTask) Execute(ctx context.Context) (err error) { dt.req.GetCollectionName(), dt.partitionID, dt.req.GetPartitionName(), dt.req.GetDbName(), + dt.req.Namespace, + collectionSchema, ) if err != nil { return err @@ -40,13 +54,7 @@ func (dt *deleteTask) Execute(ctx context.Context) (err error) { var ez *message.CipherConfig if hookutil.IsClusterEncryptionEnabled() { - schema, err := globalMetaCache.GetCollectionSchema(ctx, dt.req.GetDbName(), dt.req.GetCollectionName()) - if err != nil { - mlog.Warn(ctx, "get collection schema from global meta cache failed", mlog.String("collectionName", dt.req.GetCollectionName()), mlog.Err(err)) - return err - } - - ez = hookutil.GetEzByCollProperties(schema.GetProperties(), dt.collectionID).AsMessageConfig() + ez = hookutil.GetEzByCollProperties(collectionSchema.GetProperties(), dt.collectionID).AsMessageConfig() } var msgs []message.MutableMessage diff --git a/internal/proxy/task_delete_test.go b/internal/proxy/task_delete_test.go index 1082617474..ee3a88f80e 100644 --- a/internal/proxy/task_delete_test.go +++ b/internal/proxy/task_delete_test.go @@ -138,6 +138,19 @@ func TestDeleteTask_GetChannels(t *testing.T) { assert.ElementsMatch(t, channels, dt.pChannels) } +func TestDeleteTask_PreExecuteSkipsNamespaceValidationWhenUnset(t *testing.T) { + cache := globalMetaCache + globalMetaCache = nil + defer func() { + globalMetaCache = cache + }() + + dt := deleteTask{ + req: &milvuspb.DeleteRequest{}, + } + assert.NoError(t, dt.PreExecute(context.Background())) +} + func TestDeleteTask_Execute(t *testing.T) { collectionName := "test_delete" collectionID := int64(111) @@ -334,6 +347,38 @@ func (s *DeleteRunnerSuite) TestInitSuccess() { s.Require().Equal(0, len(dr.partitionIDs)) }) + s.Run("namespace assigns partition", func() { + mockChMgr := NewMockChannelsMgr(s.T()) + namespace := "ns-1" + partitionNames := []string{"_default_0", "_default_1"} + partitionIDs := map[string]int64{"_default_0": 100, "_default_1": 101} + schema := namespaceEnabledSchema( + &schemapb.FieldSchema{FieldID: common.StartOfUserFieldID, Name: "pk", IsPrimaryKey: true, DataType: schemapb.DataType_Int64}, + &schemapb.FieldSchema{FieldID: common.StartOfUserFieldID + 1, Name: "non_pk", DataType: schemapb.DataType_Int64}, + ) + s.schema = newSchemaInfo(schema) + dr := deleteRunner{ + req: &milvuspb.DeleteRequest{ + CollectionName: s.collectionName, + Expr: "pk == 1", + Namespace: &namespace, + }, + chMgr: mockChMgr, + } + s.mockCache.EXPECT().GetDatabaseInfo(mock.Anything, mock.Anything).Return(&databaseInfo{dbID: 0}, nil) + s.mockCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(s.collectionID, nil) + s.mockCache.EXPECT().GetCollectionInfo(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&collectionInfo{}, nil) + s.mockCache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(s.schema, nil).Twice() + s.mockCache.EXPECT().GetPartitionsIndex(mock.Anything, mock.Anything, mock.Anything).Return(partitionNames, nil) + s.mockCache.EXPECT().GetPartitions(mock.Anything, mock.Anything, mock.Anything).Return(partitionIDs, nil) + mockChMgr.EXPECT().getVChannels(mock.Anything).Return([]string{"vchan1"}, nil) + + globalMetaCache = s.mockCache + s.NoError(dr.Init(context.Background())) + + s.Equal([]int64{expectedNamespacePartitionID(namespace, partitionNames, partitionIDs)}, dr.partitionIDs) + }) + s.Run("pk == 1, no partition name", func() { mockChMgr := NewMockChannelsMgr(s.T()) dr := deleteRunner{ @@ -471,6 +516,37 @@ func (s *DeleteRunnerSuite) TestInitFailure() { s.Error(dr.Init(context.Background())) }) + s.Run("namespace disabled", func() { + mockChMgr := NewMockChannelsMgr(s.T()) + namespace := "ns-1" + dr := deleteRunner{ + req: &milvuspb.DeleteRequest{ + CollectionName: s.collectionName, + Expr: "pk == 1", + Namespace: &namespace, + }, + chMgr: mockChMgr, + } + schema := &schemapb.CollectionSchema{ + Name: s.collectionName, + Fields: []*schemapb.FieldSchema{ + { + FieldID: common.StartOfUserFieldID, + Name: "pk", + IsPrimaryKey: true, + DataType: schemapb.DataType_Int64, + }, + }, + } + s.schema = newSchemaInfo(schema) + s.mockCache.EXPECT().GetDatabaseInfo(mock.Anything, mock.Anything).Return(&databaseInfo{dbID: 0}, nil) + s.mockCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(s.collectionID, nil) + s.mockCache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(s.schema, nil) + + globalMetaCache = s.mockCache + s.Error(dr.Init(context.Background())) + }) + s.Run("fail to get database info", func() { dr := deleteRunner{ req: &milvuspb.DeleteRequest{ diff --git a/internal/proxy/task_insert_streaming.go b/internal/proxy/task_insert_streaming.go index 48465d3989..045ab7c127 100644 --- a/internal/proxy/task_insert_streaming.go +++ b/internal/proxy/task_insert_streaming.go @@ -65,7 +65,7 @@ func (it *insertTask) Execute(ctx context.Context) error { if it.partitionKeys == nil { msgs, err = repackInsertDataForStreamingService(it.TraceCtx(), channelNames, it.insertMsg, it.result, ez, it.schemaVersion) } else { - msgs, err = repackInsertDataWithPartitionKeyForStreamingService(it.TraceCtx(), channelNames, it.insertMsg, it.result, it.partitionKeys, ez, it.schemaVersion) + msgs, err = repackInsertDataWithPartitionKeyForStreamingService(it.TraceCtx(), channelNames, it.insertMsg, it.result, it.partitionKeys, ez, it.schema, it.schemaVersion) } if err != nil { mlog.Warn(ctx, "assign segmentID and repack insert data failed", mlog.Err(err)) @@ -146,11 +146,18 @@ func repackInsertDataWithPartitionKeyForStreamingService( result *milvuspb.MutationResult, partitionKeys *schemapb.FieldData, ez *message.CipherConfig, + schema *schemapb.CollectionSchema, schemaVersion int32, ) ([]message.MutableMessage, error) { messages := make([]message.MutableMessage, 0) - channel2RowOffsets, err := assignChannelsByPK(result.IDs, channelNames, insertMsg) + var channel2RowOffsets map[string][]int + var err error + if namespacePartitionKeyModeEnabled(schema) && insertMsg.Namespace != nil { + channel2RowOffsets, err = assignChannelsByNamespace(*insertMsg.Namespace, channelNames, insertMsg) + } else { + channel2RowOffsets, err = assignChannelsByPK(result.IDs, channelNames, insertMsg) + } if err != nil { return nil, err } diff --git a/internal/proxy/task_query.go b/internal/proxy/task_query.go index d72742621f..2e79e49f39 100644 --- a/internal/proxy/task_query.go +++ b/internal/proxy/task_query.go @@ -87,6 +87,7 @@ type queryTask struct { resolvedTimezoneStr string storageCost segcore.StorageCost aggregationFieldMap *agg.AggregationFieldMap + chMgr channelsMgr } func (t *queryTask) getQueryLabel() string { @@ -805,7 +806,14 @@ func (t *queryTask) PreExecute(ctx context.Context) error { // convert partition names only when requery is false if !t.reQuery { partitionNames := t.request.GetPartitionNames() - if t.partitionKeyMode { + if namespacePartitionKeyMode(t.schema.CollectionSchema) && t.request.Namespace != nil { + hashedPartitionNames, err := assignNamespacePartitionKey(ctx, t.request.GetDbName(), t.request.CollectionName, t.request.Namespace) + if err != nil { + return err + } + + partitionNames = append(partitionNames, hashedPartitionNames...) + } else if t.partitionKeyMode { expr, err := exprutil.ParseExprFromPlan(t.plan) if err != nil { return err @@ -919,6 +927,34 @@ func (t *queryTask) Execute(ctx context.Context) error { mlog.String("requestType", t.getQueryLabel())) t.resultBuf = typeutil.NewConcurrentSet[*internalpb.RetrieveResults]() + if namespacePartitionKeyModeEnabled(t.schema.CollectionSchema) && t.request.Namespace != nil { + channelNames, err := t.chMgr.getVChannels(t.CollectionID) + if err != nil { + log.Warn(ctx, "get vChannels failed", mlog.Int64("collectionID", t.CollectionID), mlog.Err(err)) + return err + } + channelName, ok, err := namespaceShardingChannel(t.schema.CollectionSchema, t.request.Namespace, channelNames) + if err != nil { + return err + } + if ok { + if err := t.lb.ExecuteWithRetry(ctx, shardclient.ChannelWorkload{ + Db: t.request.GetDbName(), + CollectionName: t.collectionName, + CollectionID: t.CollectionID, + Channel: channelName, + Nq: 1, + Exec: t.queryShard, + PreferredNodeID: preferredNodeForChannel(t.preferredNodes, channelName), + }); err != nil { + log.Warn(ctx, "fail to execute query", mlog.Err(err)) + return errors.Wrap(err, "failed to query") + } + + log.Debug(ctx, "Query Execute done.") + return nil + } + } err := t.lb.Execute(ctx, shardclient.CollectionWorkLoad{ Db: t.request.GetDbName(), CollectionID: t.CollectionID, diff --git a/internal/proxy/task_query_namespace_test.go b/internal/proxy/task_query_namespace_test.go index 1cae91f54b..f0b6910bbd 100644 --- a/internal/proxy/task_query_namespace_test.go +++ b/internal/proxy/task_query_namespace_test.go @@ -11,11 +11,46 @@ import ( "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/internal/parser/planparserv2" + "github.com/milvus-io/milvus/pkg/v3/common" "github.com/milvus-io/milvus/pkg/v3/proto/internalpb" "github.com/milvus-io/milvus/pkg/v3/proto/planpb" "github.com/milvus-io/milvus/pkg/v3/util/merr" + "github.com/milvus-io/milvus/pkg/v3/util/typeutil" ) +func namespaceEnabledSchema(fields ...*schemapb.FieldSchema) *schemapb.CollectionSchema { + return &schemapb.CollectionSchema{ + Name: "test_collection", + Fields: append(fields, &schemapb.FieldSchema{ + FieldID: 999, + Name: common.NamespaceFieldName, + IsPartitionKey: true, + DataType: schemapb.DataType_VarChar, + }), + EnableNamespace: true, + } +} + +func nonPartitionKeyPredicate(fieldID int64, dataType schemapb.DataType) *planpb.Expr { + return &planpb.Expr{ + Expr: &planpb.Expr_UnaryRangeExpr{ + UnaryRangeExpr: &planpb.UnaryRangeExpr{ + ColumnInfo: &planpb.ColumnInfo{ + FieldId: fieldID, + DataType: dataType, + }, + Op: planpb.OpType_GreaterThan, + Value: &planpb.GenericValue{Val: &planpb.GenericValue_Int64Val{Int64Val: 0}}, + }, + }, + } +} + +func expectedNamespacePartitionID(namespace string, partitionNames []string, partitionIDs map[string]int64) int64 { + idx := typeutil.HashString2Uint32(namespace) % uint32(len(partitionNames)) + return partitionIDs[partitionNames[idx]] +} + func TestQueryTask_PlanNamespace_AfterPreExecute(t *testing.T) { mockey.PatchConvey("TestQueryTask_PlanNamespace_AfterPreExecute", t, func() { // Setup global meta cache and common mocks @@ -77,3 +112,53 @@ func TestQueryTask_PlanNamespace_AfterPreExecute(t *testing.T) { assert.Equal(t, *task.request.Namespace, *capturedPlan.Namespace) }) } + +func TestQueryTask_NamespaceSetsPartitionIDs(t *testing.T) { + mockey.PatchConvey("TestQueryTask_NamespaceSetsPartitionIDs", t, func() { + globalMetaCache = &MetaCache{} + + partitionNames := []string{"_default_0", "_default_1"} + partitionIDs := map[string]int64{"_default_0": 101, "_default_1": 102} + namespaces := []string{"ns-0", "ns-1", "ns-2", "ns-3", "ns-4", "ns-5", "ns-6", "ns-7"} + schema := namespaceEnabledSchema( + &schemapb.FieldSchema{FieldID: 100, Name: "id", IsPrimaryKey: true, DataType: schemapb.DataType_Int64}, + &schemapb.FieldSchema{FieldID: 101, Name: "value", DataType: schemapb.DataType_Int64}, + ) + + mockey.Mock((*MetaCache).GetCollectionID).Return(int64(1001), nil).Build() + mockey.Mock((*MetaCache).GetCollectionInfo).Return(&collectionInfo{updateTimestamp: 12345, consistencyLevel: commonpb.ConsistencyLevel_Strong}, nil).Build() + mockey.Mock((*MetaCache).GetCollectionSchema).Return(newSchemaInfo(schema), nil).Build() + mockey.Mock((*MetaCache).GetPartitionsIndex).Return(partitionNames, nil).Build() + mockey.Mock((*MetaCache).GetPartitions).Return(partitionIDs, nil).Build() + mockey.Mock(validatePartitionTag).Return(nil).Build() + mockey.Mock(isIgnoreGrowing).Return(false, nil).Build() + + mockey.Mock((*queryTask).createPlanArgs).To(func(q *queryTask, ctx context.Context, visitorArgs *planparserv2.ParserVisitorArgs) error { + q.plan = &planpb.PlanNode{ + Node: &planpb.PlanNode_Query{ + Query: &planpb.QueryPlanNode{ + Predicates: nonPartitionKeyPredicate(101, schemapb.DataType_Int64), + }, + }, + } + q.translatedOutputFields = []string{"id"} + q.userOutputFields = []string{"id"} + return nil + }).Build() + + for _, ns := range namespaces { + namespace := ns + task := &queryTask{ + Condition: NewTaskCondition(context.Background()), + RetrieveRequest: &internalpb.RetrieveRequest{QueryLabel: "query", Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Retrieve}}, + ctx: context.Background(), + request: &milvuspb.QueryRequest{CollectionName: "test_collection", Expr: "value > 0", Namespace: &namespace}, + result: &milvuspb.QueryResults{Status: merr.Success()}, + } + + err := task.PreExecute(context.Background()) + assert.NoError(t, err) + assert.Equal(t, []int64{expectedNamespacePartitionID(namespace, partitionNames, partitionIDs)}, task.GetPartitionIDs()) + } + }) +} diff --git a/internal/proxy/task_search.go b/internal/proxy/task_search.go index 7092e7791e..a94af0e237 100644 --- a/internal/proxy/task_search.go +++ b/internal/proxy/task_search.go @@ -125,6 +125,8 @@ type searchTask struct { hybridSubSearchInfos []hybridSubSearchInfo hybridElementLevel bool + + chMgr channelsMgr } func (t *searchTask) CanSkipAllocTimestamp() bool { @@ -1179,6 +1181,23 @@ func (t *searchTask) tryGeneratePlan(params []*commonpb.KeyValuePair, dsl string } func (t *searchTask) tryParsePartitionIDsFromPlan(plan *planpb.PlanNode) ([]int64, error) { + if namespacePartitionKeyMode(t.schema.CollectionSchema) && t.request.Namespace != nil { + hashedPartitionNames, err := assignNamespacePartitionKey(t.ctx, t.request.GetDbName(), t.collectionName, t.request.Namespace) + if err != nil { + mlog.Warn(t.ctx, "failed to assign namespace partition key", mlog.Err(err)) + return nil, err + } + if len(hashedPartitionNames) > 0 { + PartitionIDs, err2 := getPartitionIDs(t.ctx, t.request.GetDbName(), t.collectionName, hashedPartitionNames) + if err2 != nil { + mlog.Warn(t.ctx, "failed to get namespace partition ids", mlog.Err(err2)) + return nil, err2 + } + return PartitionIDs, nil + } + return nil, nil + } + expr, err := exprutil.ParseExprFromPlan(plan) if err != nil { mlog.Warn(t.ctx, "failed to parse expr", mlog.Err(err)) @@ -1212,6 +1231,36 @@ func (t *searchTask) Execute(ctx context.Context) error { defer tr.CtxElapse(ctx, "done") t.queryChannelsNode = typeutil.NewConcurrentMap[string, int64]() + if namespacePartitionKeyModeEnabled(t.schema.CollectionSchema) && t.request.Namespace != nil { + channelNames, err := t.chMgr.getVChannels(t.CollectionID) + if err != nil { + log.Warn(ctx, "get vChannels failed", mlog.Int64("collectionID", t.CollectionID), mlog.Err(err)) + return err + } + channelName, ok, err := namespaceShardingChannel(t.schema.CollectionSchema, t.request.Namespace, channelNames) + if err != nil { + return err + } + if ok { + if err := t.lb.ExecuteWithRetry(ctx, shardclient.ChannelWorkload{ + Db: t.request.GetDbName(), + CollectionName: t.collectionName, + CollectionID: t.CollectionID, + Channel: channelName, + Nq: t.Nq, + Exec: t.searchShard, + PreferredNodeID: preferredNodeFromConcurrentMap(t.queryChannelsNode, channelName), + }); err != nil { + log.Warn(ctx, "search execute failed", mlog.Err(err)) + return errors.Wrap(err, "failed to search") + } + + log.Debug(ctx, "Search Execute done.", + mlog.Int64("collection", t.GetCollectionID()), + mlog.Int64s("partitionIDs", t.GetPartitionIDs())) + return nil + } + } err := t.lb.Execute(ctx, shardclient.CollectionWorkLoad{ Db: t.request.GetDbName(), CollectionID: t.CollectionID, diff --git a/internal/proxy/task_search_namespace_test.go b/internal/proxy/task_search_namespace_test.go index b1fd575ce4..76c6b3b927 100644 --- a/internal/proxy/task_search_namespace_test.go +++ b/internal/proxy/task_search_namespace_test.go @@ -70,6 +70,55 @@ func TestSearchTask_PlanNamespace_AfterPreExecute(t *testing.T) { }) } +func TestSearchTask_NamespaceSetsPartitionIDs(t *testing.T) { + mockey.PatchConvey("TestSearchTask_NamespaceSetsPartitionIDs", t, func() { + globalMetaCache = &MetaCache{} + + partitionNames := []string{"_default_0", "_default_1"} + partitionIDs := map[string]int64{"_default_0": 101, "_default_1": 102} + namespaces := []string{"ns-0", "ns-1", "ns-2", "ns-3", "ns-4", "ns-5", "ns-6", "ns-7"} + schema := namespaceEnabledSchema( + &schemapb.FieldSchema{FieldID: 100, Name: "id", IsPrimaryKey: true, DataType: schemapb.DataType_Int64}, + &schemapb.FieldSchema{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}}, + &schemapb.FieldSchema{FieldID: 102, Name: "value", DataType: schemapb.DataType_Int64}, + ) + + mockey.Mock((*MetaCache).GetCollectionID).Return(int64(1001), nil).Build() + mockey.Mock((*MetaCache).GetCollectionInfo).Return(&collectionInfo{updateTimestamp: 12345, consistencyLevel: commonpb.ConsistencyLevel_Strong}, nil).Build() + mockey.Mock((*MetaCache).GetCollectionSchema).Return(newSchemaInfo(schema), nil).Build() + mockey.Mock((*MetaCache).GetPartitionsIndex).Return(partitionNames, nil).Build() + mockey.Mock((*MetaCache).GetPartitions).Return(partitionIDs, nil).Build() + mockey.Mock(isIgnoreGrowing).Return(false, nil).Build() + mockey.Mock((*searchTask).checkNq).Return(int64(1), nil).Build() + mockey.Mock((*searchTask).tryGeneratePlan).To(func(_ *searchTask, _ []*commonpb.KeyValuePair, _ string, _ map[string]*schemapb.TemplateValue) (*planpb.PlanNode, *planpb.QueryInfo, int64, bool, []OrderByField, internalpb.SearchType, error) { + plan := &planpb.PlanNode{ + Node: &planpb.PlanNode_VectorAnns{ + VectorAnns: &planpb.VectorANNS{ + Predicates: nonPartitionKeyPredicate(102, schemapb.DataType_Int64), + }, + }, + } + qi := &planpb.QueryInfo{Topk: 10, MetricType: "L2", QueryFieldId: 101, GroupByFieldId: -1} + return plan, qi, 0, false, nil, internalpb.SearchType_DEFAULT, nil + }).Build() + + for _, ns := range namespaces { + namespace := ns + task := &searchTask{ + Condition: NewTaskCondition(context.Background()), + SearchRequest: &internalpb.SearchRequest{Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_Search}}, + ctx: context.Background(), + request: &milvuspb.SearchRequest{CollectionName: "test_collection", Namespace: &namespace}, + result: &milvuspb.SearchResults{Status: merr.Success()}, + } + + err := task.PreExecute(context.Background()) + assert.NoError(t, err) + assert.Equal(t, []int64{expectedNamespacePartitionID(namespace, partitionNames, partitionIDs)}, task.GetPartitionIDs()) + } + }) +} + func TestSearchTask_RequeryPlanNamespace(t *testing.T) { mockey.PatchConvey("TestSearchTask_RequeryPlanNamespace", t, func() { // Minimal searchTask with schema and namespace diff --git a/internal/proxy/task_upsert.go b/internal/proxy/task_upsert.go index f42837384c..8ac0c64218 100644 --- a/internal/proxy/task_upsert.go +++ b/internal/proxy/task_upsert.go @@ -222,6 +222,7 @@ func retrieveByPKs(ctx context.Context, t *upsertTask, ids *schemapb.IDs, output mixCoord: t.node.(*Proxy).mixCoord, lb: t.node.(*Proxy).lbPolicy, shardclientMgr: t.node.(*Proxy).shardMgr, + chMgr: t.node.(*Proxy).chMgr, } ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-Upsert-retrieveByPKs") diff --git a/internal/proxy/task_upsert_streaming.go b/internal/proxy/task_upsert_streaming.go index a07f892eea..5e21fd785e 100644 --- a/internal/proxy/task_upsert_streaming.go +++ b/internal/proxy/task_upsert_streaming.go @@ -88,7 +88,7 @@ func (ut *upsertTask) packInsertMessage(ctx context.Context, ez *message.CipherC if ut.partitionKeys == nil { msgs, err = repackInsertDataForStreamingService(ut.TraceCtx(), channelNames, ut.upsertMsg.InsertMsg, ut.result, ez, ut.schemaVersion) } else { - msgs, err = repackInsertDataWithPartitionKeyForStreamingService(ut.TraceCtx(), channelNames, ut.upsertMsg.InsertMsg, ut.result, ut.partitionKeys, ez, ut.schemaVersion) + msgs, err = repackInsertDataWithPartitionKeyForStreamingService(ut.TraceCtx(), channelNames, ut.upsertMsg.InsertMsg, ut.result, ut.partitionKeys, ez, ut.schema.CollectionSchema, ut.schemaVersion) } if err != nil { log.Warn(ctx, "assign segmentID and repack insert data failed", mlog.Err(err)) @@ -122,6 +122,8 @@ func (ut *upsertTask) packDeleteMessage(ctx context.Context, ez *message.CipherC ut.upsertMsg.DeleteMsg.CollectionID, ut.upsertMsg.DeleteMsg.CollectionName, ut.upsertMsg.DeleteMsg.PartitionID, ut.upsertMsg.DeleteMsg.PartitionName, ut.req.GetDbName(), + ut.req.Namespace, + ut.schema.CollectionSchema, ) if err != nil { return nil, err diff --git a/internal/proxy/util.go b/internal/proxy/util.go index 022fc5414a..bcb545a922 100644 --- a/internal/proxy/util.go +++ b/internal/proxy/util.go @@ -2648,6 +2648,30 @@ func assignChannelsByPK(pks *schemapb.IDs, channelNames []string, insertMsg *msg return channel2RowOffsets, nil } +func assignChannelsByNamespace(namespace string, channelNames []string, insertMsg *msgstream.InsertMsg) (map[string][]int, error) { + if len(channelNames) == 0 { + return nil, merr.WrapErrServiceInternalMsg("no virtual channels available for namespace sharding") + } + channelID := typeutil.HashNamespace2Channels(namespace, channelNames) + return assignChannelsByChannel(channelID, channelNames, insertMsg), nil +} + +func assignChannelsByChannel(channelID uint32, channelNames []string, insertMsg *msgstream.InsertMsg) map[string][]int { + insertMsg.HashValues = make([]uint32, insertMsg.NumRows) + for i := range insertMsg.HashValues { + insertMsg.HashValues[i] = channelID + } + + channelName := channelNames[channelID] + channel2RowOffsets := map[string][]int{ + channelName: make([]int, 0, insertMsg.NRows()), + } + for i := range insertMsg.HashValues { + channel2RowOffsets[channelName] = append(channel2RowOffsets[channelName], i) + } + return channel2RowOffsets +} + func assignPartitionKeys(ctx context.Context, dbName string, collName string, keys []*planpb.GenericValue) ([]string, error) { partitionNames, err := globalMetaCache.GetPartitionsIndex(ctx, dbName, collName) if err != nil { @@ -2668,6 +2692,16 @@ func assignPartitionKeys(ctx context.Context, dbName string, collName string, ke return hashedPartitionNames, err } +func assignNamespacePartitionKey(ctx context.Context, dbName string, collName string, namespace *string) ([]string, error) { + if namespace == nil { + return nil, nil + } + + return assignPartitionKeys(ctx, dbName, collName, []*planpb.GenericValue{ + {Val: &planpb.GenericValue_StringVal{StringVal: *namespace}}, + }) +} + func ErrWithLog(logger *mlog.Logger, msg string, err error) error { wrapErr := errors.Wrap(err, msg) if logger != nil { @@ -2745,8 +2779,64 @@ func checkDynamicFieldDataForPartialUpdate(schema *schemapb.CollectionSchema, in return doCheckDynamicFieldData(schema, insertMsg, true) } +func namespaceShardingEnabled(schema *schemapb.CollectionSchema) bool { + if schema == nil || !schema.GetEnableNamespace() { + return false + } + enabled, err := common.IsNamespaceShardingEnabled(schema.GetProperties()...) + return err == nil && enabled +} + +func namespacePartitionKeyMode(schema *schemapb.CollectionSchema) bool { + return schema != nil && schema.GetEnableNamespace() && common.IsNamespaceModePartitionKey(schema.GetProperties()...) +} + +func namespacePartitionKeyModeEnabled(schema *schemapb.CollectionSchema) bool { + return namespaceShardingEnabled(schema) && namespacePartitionKeyMode(schema) +} + func namespacePartitionModeEnabled(schema *schemapb.CollectionSchema) bool { - return schema.GetEnableNamespace() && common.IsNamespaceModePartition(schema.GetProperties()...) + return schema != nil && schema.GetEnableNamespace() && common.IsNamespaceModePartition(schema.GetProperties()...) +} + +func namespaceShardingChannelID(schema *schemapb.CollectionSchema, namespace *string, channelNames []string) (uint32, bool, error) { + if namespace == nil || !namespacePartitionKeyModeEnabled(schema) { + return 0, false, nil + } + if len(channelNames) == 0 { + return 0, false, merr.WrapErrServiceInternalMsg("no virtual channels available for namespace sharding") + } + return typeutil.HashNamespace2Channels(*namespace, channelNames), true, nil +} + +func namespaceShardingChannel(schema *schemapb.CollectionSchema, namespace *string, channelNames []string) (string, bool, error) { + channelID, ok, err := namespaceShardingChannelID(schema, namespace, channelNames) + if !ok || err != nil { + return "", ok, err + } + return channelNames[channelID], true, nil +} + +func preferredNodeForChannel(preferredNodes map[string]int64, channel string) int64 { + if preferredNodes == nil { + return 0 + } + preferredNodeID, ok := preferredNodes[channel] + if !ok { + return 0 + } + return preferredNodeID +} + +func preferredNodeFromConcurrentMap(preferredNodes *typeutil.ConcurrentMap[string, int64], channel string) int64 { + if preferredNodes == nil { + return 0 + } + preferredNodeID, ok := preferredNodes.Get(channel) + if !ok { + return 0 + } + return preferredNodeID } func resolveNamespacePartitionName(schema *schemapb.CollectionSchema, namespace *string, partitionName string) (string, bool, error) { diff --git a/pkg/util/typeutil/hash.go b/pkg/util/typeutil/hash.go index 104d401095..20b58f18e1 100644 --- a/pkg/util/typeutil/hash.go +++ b/pkg/util/typeutil/hash.go @@ -187,3 +187,9 @@ func RearrangePartitionsForPartitionKey(partitions map[string]int64) ([]string, return partitionNames, partitionIDs, nil } + +func HashNamespace2Channels(namespace string, shardNames []string) uint32 { + numShard := uint32(len(shardNames)) + hashValue := HashString2Uint32(namespace) + return hashValue % numShard +} diff --git a/scripts/stop.sh b/scripts/stop.sh index d697fe03b7..b7127d6914 100755 --- a/scripts/stop.sh +++ b/scripts/stop.sh @@ -1,3 +1,5 @@ +#!/usr/bin/env bash + # 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 @@ -14,12 +16,34 @@ # See the License for the specific language governing permissions and # limitations under the License. -echo "Stopping milvus..." -PROCESS=$(ps -e | grep milvus | grep -v grep | awk '{print $1}') -if [ -z "$PROCESS" ]; then - echo "No milvus process" - exit 0 -fi -kill -15 $PROCESS -echo "Milvus stopped" +function filter_milvus_process() { + awk ' + { + pid = $1 + $1 = "" + sub(/^[[:space:]]+/, "", $0) + # Match Milvus commands, not tools/editors that mention the repo name. + if ($0 ~ /^([^[:space:]]*\/)?milvus[[:space:]]+run[[:space:]]/) { + print pid + } + }' +} +function get_milvus_process() { + ps -e -o pid= -o args= | filter_milvus_process +} + +function main() { + echo "Stopping milvus..." + PROCESS=$(get_milvus_process) + if [ -z "$PROCESS" ]; then + echo "No milvus process" + exit 0 + fi + kill -15 $PROCESS + echo "Milvus stopped" +} + +if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then + main "$@" +fi diff --git a/tests/python_client/milvus_client/test_milvus_client_query.py b/tests/python_client/milvus_client/test_milvus_client_query.py index e984210570..b154216cf6 100644 --- a/tests/python_client/milvus_client/test_milvus_client_query.py +++ b/tests/python_client/milvus_client/test_milvus_client_query.py @@ -805,6 +805,68 @@ class TestMilvusClientQueryValid(TestMilvusClientV2Base): ) self.drop_collection(client, collection_name) + @pytest.mark.tags(CaseLabel.L1) + def test_milvus_client_query_with_multiple_namespaces(self): + """ + target: test query with namespace enabled + method: create namespace-enabled collection with multiple shards, insert rows into multiple namespaces, then query by namespace + expected: query only returns rows from the requested namespace + """ + client = self._client() + collection_name = cf.gen_collection_name_by_testcase_name() + shards_num = 4 + self.create_collection( + client, + collection_name, + default_dim, + consistency_level="Strong", + enable_namespace=True, + shards_num=shards_num, + ) + + desc = self.describe_collection(client, collection_name)[0] + assert desc["enable_namespace"] is True + assert desc["num_shards"] == shards_num + + rng = np.random.default_rng(seed=19530) + namespaces = [f"namespace_{i}" for i in range(8)] + nb_per_namespace = 10 + ids_by_namespace = { + namespace: [idx * 1000 + i for i in range(nb_per_namespace)] + for idx, namespace in enumerate(namespaces) + } + for namespace, ids in ids_by_namespace.items(): + rows = [ + { + default_primary_key_field_name: pk, + default_vector_field_name: list(rng.random((1, default_dim))[0]), + } + for pk in ids + ] + self.insert(client, collection_name, rows, namespace=namespace) + + for namespace, expected_ids in ids_by_namespace.items(): + res = self.query( + client, + collection_name, + filter=default_search_exp, + output_fields=[default_primary_key_field_name], + namespace=namespace, + )[0] + assert {r[default_primary_key_field_name] for r in res} == set(expected_ids) + + first_id_per_namespace = [ids[0] for ids in ids_by_namespace.values()] + for namespace, expected_ids in ids_by_namespace.items(): + filtered = self.query( + client, + collection_name, + filter=f"{default_primary_key_field_name} in {first_id_per_namespace}", + output_fields=[default_primary_key_field_name], + namespace=namespace, + )[0] + assert {r[default_primary_key_field_name] for r in filtered} == {expected_ids[0]} + self.drop_collection(client, collection_name) + @pytest.mark.tags(CaseLabel.L1) def test_milvus_client_query_output_fields(self): """