Improve file naming and schema creation usage (#458)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
milvus-io · May 22, 2023 · f6daf36 · f6daf36
1 parent b619b53
commit f6daf36
Show file tree

Hide file tree

Showing 24 changed files with 97 additions and 205 deletions.
diff --git a/client/client_grpc_admin.go → client/admin.go b/client/client_grpc_admin.go → client/admin.go
@@ -17,6 +17,7 @@ import (
 	server "github.com/milvus-io/milvus-proto/go-api/milvuspb"
 )
 
+// GetVersion returns milvus server version information.
 func (c *GrpcClient) GetVersion(ctx context.Context) (string, error) {
 	if c.Service == nil {
 		return "", ErrClientNotReady

diff --git a/client/client_grpc_alias.go → client/alias.go b/client/client_grpc_alias.go → client/alias.go
diff --git a/client/client_grpc_alias_test.go → client/alias_test.go b/client/client_grpc_alias_test.go → client/alias_test.go
diff --git a/client/client_grpc_authentication.go → client/authentication.go b/client/client_grpc_authentication.go → client/authentication.go
diff --git a/client/client_grpc_authentication_test.go → client/authentication_test.go b/client/client_grpc_authentication_test.go → client/authentication_test.go
diff --git a/client/client_test.go b/client/client_test.go
@@ -47,50 +47,15 @@ const (
 )
 
 func defaultSchema() *entity.Schema {
-	return &entity.Schema{
-		CollectionName: testCollectionName,
-		AutoID:         false,
-		Fields: []*entity.Field{
-			{
-				Name:       testPrimaryField,
-				DataType:   entity.FieldTypeInt64,
-				PrimaryKey: true,
-				AutoID:     true,
-			},
-			{
-				Name:     testVectorField,
-				DataType: entity.FieldTypeFloatVector,
-				TypeParams: map[string]string{
-					entity.TypeParamDim: fmt.Sprintf("%d", testVectorDim),
-				},
-			},
-		},
-	}
+	return entity.NewSchema().WithName(testCollectionName).WithAutoID(false).
+		WithField(entity.NewField().WithName(testPrimaryField).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(true)).
+		WithField(entity.NewField().WithName(testVectorField).WithDataType(entity.FieldTypeFloatVector).WithDim(testVectorDim))
 }
 
 func varCharSchema() *entity.Schema {
-	return &entity.Schema{
-		CollectionName: testCollectionName,
-		AutoID:         false,
-		Fields: []*entity.Field{
-			{
-				Name:       "varchar",
-				DataType:   entity.FieldTypeVarChar,
-				PrimaryKey: true,
-				AutoID:     false,
-				TypeParams: map[string]string{
-					entity.TypeParamMaxLength: fmt.Sprintf("%d", 100),
-				},
-			},
-			{
-				Name:     testVectorField,
-				DataType: entity.FieldTypeFloatVector,
-				TypeParams: map[string]string{
-					entity.TypeParamDim: fmt.Sprintf("%d", testVectorDim),
-				},
-			},
-		},
-	}
+	return entity.NewSchema().WithName(testCollectionName).WithAutoID(false).
+		WithField(entity.NewField().WithName("varchar").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(false).WithMaxLength(100)).
+		WithField(entity.NewField().WithName(testVectorField).WithDataType(entity.FieldTypeFloatVector).WithDim(testVectorDim))
 }
 
 var _ entity.Row = &defaultRow{}

diff --git a/client/collection.go b/client/collection.go
@@ -20,43 +20,11 @@ import (
 
 	"github.com/golang/protobuf/proto"
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
-	"google.golang.org/grpc"
 
 	common "github.com/milvus-io/milvus-proto/go-api/commonpb"
 	server "github.com/milvus-io/milvus-proto/go-api/milvuspb"
 )
 
-// GrpcClient, uses default grpc Service definition to connect with Milvus2.0
-type GrpcClient struct {
-	Conn    *grpc.ClientConn           // grpc connection instance
-	Service server.MilvusServiceClient // Service client stub
-}
-
-// connect connect to Service
-func (c *GrpcClient) connect(ctx context.Context, addr string, opts ...grpc.DialOption) error {
-	if addr == "" {
-		return fmt.Errorf("address is empty")
-	}
-	conn, err := grpc.DialContext(ctx, addr, opts...)
-	if err != nil {
-		return err
-	}
-
-	c.Conn = conn
-	c.Service = server.NewMilvusServiceClient(c.Conn)
-	return nil
-}
-
-// Close close the connection
-func (c *GrpcClient) Close() error {
-	if c.Conn != nil {
-		err := c.Conn.Close()
-		c.Conn = nil
-		return err
-	}
-	return nil
-}
-
 // handles response status
 // if status is nil returns ErrStatusNil
 // if status.ErrorCode is common.ErrorCode_Success, returns nil
@@ -222,7 +190,7 @@ func (c *GrpcClient) DescribeCollection(ctx context.Context, collName string) (*
 	collection := &entity.Collection{
 		ID:               resp.GetCollectionID(),
 		Name:             collName,
-		Schema:           (&entity.Schema{}).ReadProto(resp.GetSchema()),
+		Schema:           entity.NewSchema().ReadProto(resp.GetSchema()),
 		PhysicalChannels: resp.GetPhysicalChannelNames(),
 		VirtualChannels:  resp.GetVirtualChannelNames(),
 		ConsistencyLevel: entity.ConsistencyLevel(resp.ConsistencyLevel),

diff --git a/client/collection_test.go b/client/collection_test.go
@@ -166,97 +166,34 @@ func (s *CollectionSuite) TestCreateCollection() {
 		}
 		cases := []testCase{
 			{
-				name: "empty_fields",
-				schema: &entity.Schema{
-					CollectionName: testCollectionName,
-					Fields:         []*entity.Field{},
-				},
+				name:   "empty_fields",
+				schema: entity.NewSchema().WithName(testCollectionName),
 			},
 			{
 				name: "empty_collection_name",
-				schema: &entity.Schema{
-					CollectionName: "",
-					Fields: []*entity.Field{
-						{
-							Name:       "int64",
-							DataType:   entity.FieldTypeInt64,
-							PrimaryKey: true,
-						},
-						{
-							Name:       "vector",
-							DataType:   entity.FieldTypeFloatVector,
-							TypeParams: map[string]string{entity.TypeParamDim: "128"},
-						},
-					},
-				},
+				schema: entity.NewSchema().
+					WithField(entity.NewField().WithName("int64").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
+					WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128)),
 			},
-
 			{
 				name: "multiple primary key",
-				schema: &entity.Schema{
-
-					CollectionName: testCollectionName,
-					Fields: []*entity.Field{
-						{
-							Name:       "int64",
-							DataType:   entity.FieldTypeInt64,
-							PrimaryKey: true,
-						},
-						{
-							Name:       "int64_2",
-							DataType:   entity.FieldTypeInt64,
-							PrimaryKey: true,
-						},
-						{
-							Name:       "vector",
-							DataType:   entity.FieldTypeFloatVector,
-							TypeParams: map[string]string{entity.TypeParamDim: "128"},
-						},
-					},
-				},
+				schema: entity.NewSchema().WithName(testCollectionName).
+					WithField(entity.NewField().WithName("int64").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
+					WithField(entity.NewField().WithName("int64_2").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
+					WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128)),
 			},
 			{
 				name: "multiple auto id",
-				schema: &entity.Schema{
-					CollectionName: testCollectionName,
-					Fields: []*entity.Field{
-						{
-							Name:       "int64",
-							DataType:   entity.FieldTypeInt64,
-							PrimaryKey: true,
-							AutoID:     true,
-						},
-						{
-							Name:       "int64_2",
-							DataType:   entity.FieldTypeInt64,
-							PrimaryKey: false,
-							AutoID:     true,
-						},
-						{
-							Name:       "vector",
-							DataType:   entity.FieldTypeFloatVector,
-							TypeParams: map[string]string{entity.TypeParamDim: "128"},
-						},
-					},
-				},
+				schema: entity.NewSchema().WithName(testCollectionName).
+					WithField(entity.NewField().WithName("int64").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(true)).
+					WithField(entity.NewField().WithName("int64_2").WithDataType(entity.FieldTypeInt64).WithIsAutoID(true)).
+					WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128)),
 			},
 			{
 				name: "bad_pk_type",
-				schema: &entity.Schema{
-					CollectionName: testCollectionName,
-					Fields: []*entity.Field{
-						{
-							Name:       "float_pk",
-							DataType:   entity.FieldTypeFloat,
-							PrimaryKey: true,
-						},
-						{
-							Name:       "vector",
-							DataType:   entity.FieldTypeFloatVector,
-							TypeParams: map[string]string{entity.TypeParamDim: "128"},
-						},
-					},
-				},
+				schema: entity.NewSchema().
+					WithField(entity.NewField().WithName("int64").WithDataType(entity.FieldTypeDouble).WithIsPrimaryKey(true)).
+					WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128)),
 			},
 		}
 

diff --git a/client/connect.go b/client/connect.go
@@ -0,0 +1,51 @@
+// Copyright (C) 2019-2021 Zilliz. All rights reserved.
+//
+// Licensed 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 client
+
+import (
+	"context"
+	"fmt"
+
+	server "github.com/milvus-io/milvus-proto/go-api/milvuspb"
+	"google.golang.org/grpc"
+)
+
+// GrpcClient, uses default grpc Service definition to connect with Milvus2.0
+type GrpcClient struct {
+	Conn    *grpc.ClientConn           // grpc connection instance
+	Service server.MilvusServiceClient // Service client stub
+}
+
+// connect connect to Service
+func (c *GrpcClient) connect(ctx context.Context, addr string, opts ...grpc.DialOption) error {
+	if addr == "" {
+		return fmt.Errorf("address is empty")
+	}
+	conn, err := grpc.DialContext(ctx, addr, opts...)
+	if err != nil {
+		return err
+	}
+
+	c.Conn = conn
+	c.Service = server.NewMilvusServiceClient(c.Conn)
+	return nil
+}
+
+// Close close the connection
+func (c *GrpcClient) Close() error {
+	if c.Conn != nil {
+		err := c.Conn.Close()
+		c.Conn = nil
+		return err
+	}
+	return nil
+}
diff --git a/client/data_test.go b/client/data_test.go
@@ -13,7 +13,6 @@ package client
 
 import (
 	"context"
-	"fmt"
 	"math/rand"
 	"testing"
 	"time"
@@ -954,56 +953,17 @@ func TestIsCollectionPrimaryKey(t *testing.T) {
 
 func TestEstRowSize(t *testing.T) {
 	// a schema contains all supported vector
-	sch := &entity.Schema{
-		CollectionName: testCollectionName,
-		AutoID:         false,
-		Fields: []*entity.Field{
-			{
-				Name:       testPrimaryField,
-				DataType:   entity.FieldTypeInt64,
-				PrimaryKey: true,
-				AutoID:     true,
-			},
-			{
-				Name:     "attr1",
-				DataType: entity.FieldTypeInt8,
-			},
-			{
-				Name:     "attr2",
-				DataType: entity.FieldTypeInt16,
-			},
-			{
-				Name:     "attr3",
-				DataType: entity.FieldTypeInt32,
-			},
-			{
-				Name:     "attr4",
-				DataType: entity.FieldTypeFloat,
-			},
-			{
-				Name:     "attr5",
-				DataType: entity.FieldTypeDouble,
-			},
-			{
-				Name:     "attr6",
-				DataType: entity.FieldTypeBool,
-			},
-			{
-				Name:     testVectorField,
-				DataType: entity.FieldTypeFloatVector,
-				TypeParams: map[string]string{
-					entity.TypeParamDim: fmt.Sprintf("%d", testVectorDim),
-				},
-			},
-			{
-				Name:     "binary_vector",
-				DataType: entity.FieldTypeBinaryVector,
-				TypeParams: map[string]string{
-					entity.TypeParamDim: fmt.Sprintf("%d", testVectorDim),
-				},
-			},
-		},
-	}
+	sch := entity.NewSchema().WithName(testCollectionName).WithAutoID(false).
+		WithField(entity.NewField().WithName(testPrimaryField).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true).WithIsAutoID(true)).
+		WithField(entity.NewField().WithName("attr1").WithDataType(entity.FieldTypeInt8)).
+		WithField(entity.NewField().WithName("attr2").WithDataType(entity.FieldTypeInt16)).
+		WithField(entity.NewField().WithName("attr3").WithDataType(entity.FieldTypeInt32)).
+		WithField(entity.NewField().WithName("attr4").WithDataType(entity.FieldTypeFloat)).
+		WithField(entity.NewField().WithName("attr5").WithDataType(entity.FieldTypeDouble)).
+		WithField(entity.NewField().WithName("attr6").WithDataType(entity.FieldTypeBool)).
+		WithField(entity.NewField().WithName("attr6").WithDataType(entity.FieldTypeBool)).
+		WithField(entity.NewField().WithName(testVectorField).WithDataType(entity.FieldTypeFloatVector).WithDim(testVectorDim)).
+		WithField(entity.NewField().WithName("binary_vector").WithDataType(entity.FieldTypeBinaryVector).WithDim(testVectorDim))
 
 	// one row
 	columnID := entity.NewColumnInt64(testPrimaryField, []int64{0})

diff --git a/client/client_grpc_index.go → client/index.go b/client/client_grpc_index.go → client/index.go
diff --git a/client/client_grpc_index_test.go → client/index_test.go b/client/client_grpc_index_test.go → client/index_test.go
diff --git a/client/client_grpc_control.go → client/maintainance.go b/client/client_grpc_control.go → client/maintainance.go
diff --git a/client/client_grpc_control_test.go → client/maintainance_test.go b/client/client_grpc_control_test.go → client/maintainance_test.go
diff --git a/client/client_grpc_options.go → client/options.go b/client/client_grpc_options.go → client/options.go
diff --git a/client/client_grpc_options_test.go → client/options_test.go b/client/client_grpc_options_test.go → client/options_test.go
diff --git a/client/client_grpc_partition.go → client/partition.go b/client/client_grpc_partition.go → client/partition.go
diff --git a/client/client_grpc_partition_test.go → client/partition_test.go b/client/client_grpc_partition_test.go → client/partition_test.go
diff --git a/client/client_grpc_rbac.go → client/rbac.go b/client/client_grpc_rbac.go → client/rbac.go
diff --git a/client/client_grpc_rbac_test.go → client/rbac_test.go b/client/client_grpc_rbac_test.go → client/rbac_test.go
diff --git a/client/client_grpc_rg.go → client/resource_group.go b/client/client_grpc_rg.go → client/resource_group.go
diff --git a/client/client_grpc_rg_test.go → client/resource_group_test.go b/client/client_grpc_rg_test.go → client/resource_group_test.go
diff --git a/entity/schema.go b/entity/schema.go
@@ -220,6 +220,14 @@ func (f *Field) WithDim(dim int64) *Field {
 	return f
 }
 
+func (f *Field) WithMaxLength(maxLen int64) *Field {
+	if f.TypeParams == nil {
+		f.TypeParams = make(map[string]string)
+	}
+	f.TypeParams[TypeParamMaxLength] = strconv.FormatInt(maxLen, 10)
+	return f
+}
+
 // ReadProto parses FieldSchema
 func (f *Field) ReadProto(p *schema.FieldSchema) *Field {
 	f.ID = p.GetFieldID()

diff --git a/entity/schema_test.go b/entity/schema_test.go
@@ -65,12 +65,14 @@ func (s *SchemaSuite) TestBasic() {
 		{
 			"test_collection",
 			NewSchema().WithName("test_collection_1").WithDescription("test_collection_1 desc").WithAutoID(false).
-				WithField(NewField().WithName("ID").WithDataType(FieldTypeInt64).WithIsPrimaryKey(true)),
+				WithField(NewField().WithName("ID").WithDataType(FieldTypeInt64).WithIsPrimaryKey(true)).
+				WithField(NewField().WithName("vector").WithDataType(FieldTypeFloatVector).WithDim(128)),
 			"ID",
 		},
 		{
 			"dynamic_schema",
 			NewSchema().WithName("dynamic_schema").WithDescription("dynamic_schema desc").WithAutoID(true).WithDynamicFieldEnabled(true).
+				WithField(NewField().WithName("ID").WithDataType(FieldTypeVarChar).WithMaxLength(256)).
 				WithField(NewField().WithName("$meta").WithIsDynamic(true)),
 			"",
 		},