Support JSON datatype

client/client_mock_test.go

@@ -9,6 +9,8 @@ import (
 	"github.com/golang/protobuf/proto"
 	common "github.com/milvus-io/milvus-proto/go-api/commonpb"
 	server "github.com/milvus-io/milvus-proto/go-api/milvuspb"
+	schema "github.com/milvus-io/milvus-proto/go-api/schemapb"
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 	"github.com/milvus-io/milvus-sdk-go/v2/mocks"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
@@ -77,17 +79,132 @@ func (s *MockSuiteBase) resetMock() {
 	}
 }
 
-func (s *MockSuiteBase) setupHasCollection(collName string) {
+func (s *MockSuiteBase) setupHasCollection(collNames ...string) {
 	s.mock.EXPECT().HasCollection(mock.Anything, mock.AnythingOfType("*milvuspb.HasCollectionRequest")).
 		Call.Return(func(ctx context.Context, req *server.HasCollectionRequest) *server.BoolResponse {
 		resp := &server.BoolResponse{Status: &common.Status{}}
+		for _, collName := range collNames {
+			if req.GetCollectionName() == collName {
+				resp.Value = true
+				break
+			}
+		}
+		return resp
+	}, nil)
+}
+
+func (s *MockSuiteBase) setupHasCollectionError(errorCode common.ErrorCode, err error) {
+	s.mock.EXPECT().HasCollection(mock.Anything, mock.AnythingOfType("*milvuspb.HasCollectionRequest")).
+		Return(&server.BoolResponse{
+			Status: &common.Status{ErrorCode: errorCode},
+		}, err)
+}
+
+func (s *MockSuiteBase) setupHasPartition(collName string, partNames ...string) {
+	s.mock.EXPECT().HasPartition(mock.Anything, mock.AnythingOfType("*milvuspb.HasPartitionRequest")).
+		Call.Return(func(ctx context.Context, req *server.HasPartitionRequest) *server.BoolResponse {
+		resp := &server.BoolResponse{Status: &common.Status{}}
 		if req.GetCollectionName() == collName {
-			resp.Value = true
+			for _, partName := range partNames {
+				if req.GetPartitionName() == partName {
+					resp.Value = true
+					break
+				}
+			}
 		}
 		return resp
 	}, nil)
 }
 
+func (s *MockSuiteBase) setupHasPartitionError(errorCode common.ErrorCode, err error) {
+	s.mock.EXPECT().HasPartition(mock.Anything, mock.AnythingOfType("*milvuspb.HasPartitionRequest")).
+		Return(&server.BoolResponse{
+			Status: &common.Status{ErrorCode: errorCode},
+		}, err)
+}
+
+func (s *MockSuiteBase) setupDescribeCollection(collName string, schema *entity.Schema) {
+	s.mock.EXPECT().DescribeCollection(mock.Anything, mock.AnythingOfType("*milvuspb.DescribeCollectionRequest")).
+		Call.Return(func(ctx context.Context, req *server.DescribeCollectionRequest) *server.DescribeCollectionResponse {
+		return &server.DescribeCollectionResponse{
+			Status: &common.Status{ErrorCode: common.ErrorCode_Success},
+			Schema: schema.ProtoMessage(),
+		}
+	}, nil)
+}
+
+func (s *MockSuiteBase) setupDescribeCollectionError(errorCode common.ErrorCode, err error) {
+	s.mock.EXPECT().DescribeCollection(mock.Anything, mock.AnythingOfType("*milvuspb.DescribeCollectionRequest")).
+		Return(&server.DescribeCollectionResponse{
+			Status: &common.Status{ErrorCode: errorCode},
+		}, err)
+}
+
+func (s *MockSuiteBase) getInt64FieldData(name string, data []int64) *schema.FieldData {
+	return &schema.FieldData{
+		Type:      schema.DataType_Int64,
+		FieldName: name,
+		Field: &schema.FieldData_Scalars{
+			Scalars: &schema.ScalarField{
+				Data: &schema.ScalarField_LongData{
+					LongData: &schema.LongArray{
+						Data: data,
+					},
+				},
+			},
+		},
+	}
+}
+
+func (s *MockSuiteBase) getVarcharFieldData(name string, data []string) *schema.FieldData {
+	return &schema.FieldData{
+		Type:      schema.DataType_VarChar,
+		FieldName: name,
+		Field: &schema.FieldData_Scalars{
+			Scalars: &schema.ScalarField{
+				Data: &schema.ScalarField_StringData{
+					StringData: &schema.StringArray{
+						Data: data,
+					},
+				},
+			},
+		},
+	}
+}
+
+func (s *MockSuiteBase) getJSONBytesFieldData(name string, data [][]byte) *schema.FieldData {
+	return &schema.FieldData{
+		Type:      schema.DataType_JSON,
+		FieldName: name,
+		Field: &schema.FieldData_Scalars{
+			Scalars: &schema.ScalarField{
+				Data: &schema.ScalarField_JsonData{
+					JsonData: &schema.JSONArray{
+						Data: data,
+					},
+				},
+			},
+		},
+	}
+}
+
+func (s *MockSuiteBase) getFloatVectorFieldData(name string, dim int64, data []float32) *schema.FieldData {
+	return &schema.FieldData{
+		Type:      schema.DataType_FloatVector,
+		FieldName: name,
+		Field: &schema.FieldData_Vectors{
+			Vectors: &schema.VectorField{
+				Dim: dim,
+				Data: &schema.VectorField_FloatVector{
+					FloatVector: &schema.FloatArray{
+						Data: data,
+					},
+				},
+			},
+		},
+	}
+}
+
 // ref https://stackoverflow.com/questions/42102496/testing-a-grpc-service
 
 var errNotImplemented = errors.New("not implemented")
@@ -788,6 +905,10 @@ func (m *MockServer) GetFlushAllState(_a0 context.Context, _a1 *server.GetFlushA
 	panic("not implemented")
 }
 
+func (m *MockServer) Connect(_ context.Context, _ *server.ConnectRequest) (*server.ConnectResponse, error) {
+	panic("not implemented") // TODO: Implement
+}
+
 func SuccessStatus() (*common.Status, error) {
 	return &common.Status{ErrorCode: common.ErrorCode_Success}, nil
 }

entity/columns_json.go

@@ -0,0 +1,92 @@
+package entity
+
+import (
+	"fmt"
+
+	"github.com/cockroachdb/errors"
+	schema "github.com/milvus-io/milvus-proto/go-api/schemapb"
+)
+
+var _ (Column) = (*ColumnJSONBytes)(nil)
+
+// ColumnJSONBytes column type for JSON.
+// all items are marshaled json bytes.
+type ColumnJSONBytes struct {
+	name   string
+	values [][]byte
+}
+
+// Name returns column name.
+func (c *ColumnJSONBytes) Name() string {
+	return c.name
+}
+
+// Type returns column FieldType.
+func (c *ColumnJSONBytes) Type() FieldType {
+	return FieldTypeJSON
+}
+
+// Len returns column values length.
+func (c *ColumnJSONBytes) Len() int {
+	return len(c.values)
+}
+
+// Get returns value at index as interface{}.
+func (c *ColumnJSONBytes) Get(idx int) (interface{}, error) {
+	if idx < 0 || idx > c.Len() {
+		return nil, errors.New("index out of range")
+	}
+	return c.values[idx], nil
+}
+
+// FieldData return column data mapped to schema.FieldData.
+func (c *ColumnJSONBytes) FieldData() *schema.FieldData {
+	fd := &schema.FieldData{
+		Type:      schema.DataType_JSON,
+		FieldName: c.name,
+	}
+
+	fd.Field = &schema.FieldData_Scalars{
+		Scalars: &schema.ScalarField{
+			Data: &schema.ScalarField_JsonData{
+				JsonData: &schema.JSONArray{
+					Data: c.values,
+				},
+			},
+		},
+	}
+
+	return fd
+}
+
+// ValueByIdx returns value of the provided index.
+func (c *ColumnJSONBytes) ValueByIdx(idx int) ([]byte, error) {
+	if idx < 0 || idx >= c.Len() {
+		return nil, errors.New("index out of range")
+	}
+	return c.values[idx], nil
+}
+
+// AppendValue append value into column.
+func (c *ColumnJSONBytes) AppendValue(i interface{}) error {
+	v, ok := i.([]byte)
+	if !ok {
+		return fmt.Errorf("invalid type, expected []byte, got %T", i)
+	}
+	c.values = append(c.values, v)
+
+	return nil
+}
+
+// Data returns column data.
+func (c *ColumnJSONBytes) Data() [][]byte {
+	return c.values
+}
+
+// NewColumnJSONBytes composes a Column with json bytes.
+func NewColumnJSONBytes(name string, values [][]byte) *ColumnJSONBytes {
+	return &ColumnJSONBytes{
+		name:   name,
+		values: values,
+	}
+}

entity/columns_json_test.go

@@ -0,0 +1,78 @@
+package entity
+
+import (
+	"fmt"
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/suite"
+)
+
+type ColumnJSONBytesSuite struct {
+	suite.Suite
+}
+
+func (s *ColumnJSONBytesSuite) SetupSuite() {
+	rand.Seed(time.Now().UnixNano())
+}
+
+func (s *ColumnJSONBytesSuite) TestAttrMethods() {
+	columnName := fmt.Sprintf("column_jsonbs_%d", rand.Int())
+	columnLen := 8 + rand.Intn(10)
+
+	v := make([][]byte, columnLen)
+	column := NewColumnJSONBytes(columnName, v)
+
+	s.Run("test_meta", func() {
+		ft := FieldTypeJSON
+		s.Equal("JSON", ft.Name())
+		s.Equal("JSON", ft.String())
+		pbName, pbType := ft.PbFieldType()
+		s.Equal("JSON", pbName)
+		s.Equal("JSON", pbType)
+	})
+
+	s.Run("test_column_attribute", func() {
+		s.Equal(columnName, column.Name())
+		s.Equal(FieldTypeJSON, column.Type())
+		s.Equal(columnLen, column.Len())
+		s.EqualValues(v, column.Data())
+	})
+
+	s.Run("test_column_field_data", func() {
+		fd := column.FieldData()
+		s.NotNil(fd)
+		s.Equal(fd.GetFieldName(), columnName)
+	})
+
+	s.Run("test_column_valuer_by_idx", func() {
+		_, err := column.ValueByIdx(-1)
+		s.Error(err)
+		_, err = column.ValueByIdx(columnLen)
+		s.Error(err)
+		for i := 0; i < columnLen; i++ {
+			v, err := column.ValueByIdx(i)
+			s.NoError(err)
+			s.Equal(column.values[i], v)
+		}
+	})
+
+	s.Run("test_append_value", func() {
+		item := make([]byte, 10)
+		err := column.AppendValue(item)
+		s.NoError(err)
+		s.Equal(columnLen+1, column.Len())
+		val, err := column.ValueByIdx(columnLen)
+		s.NoError(err)
+		s.Equal(item, val)
+
+		err = column.AppendValue(1)
+		s.Error(err)
+	})
+
+}
+
+func TestColumnJSONBytes(t *testing.T) {
+	suite.Run(t, new(ColumnJSONBytesSuite))
+}

entity/schema.go

@@ -163,6 +163,8 @@ func (t FieldType) Name() string {
 		return "String"
 	case FieldTypeVarChar:
 		return "VarChar"
+	case FieldTypeJSON:
+		return "JSON"
 	case FieldTypeBinaryVector:
 		return "BinaryVector"
 	case FieldTypeFloatVector:
@@ -193,6 +195,8 @@ func (t FieldType) String() string {
 		return "string"
 	case FieldTypeVarChar:
 		return "string"
+	case FieldTypeJSON:
+		return "JSON"
 	case FieldTypeBinaryVector:
 		return "[]byte"
 	case FieldTypeFloatVector:
@@ -223,6 +227,8 @@ func (t FieldType) PbFieldType() (string, string) {
 		return "String", "string"
 	case FieldTypeVarChar:
 		return "VarChar", "string"
+	case FieldTypeJSON:
+		return "JSON", "JSON"
 	case FieldTypeBinaryVector:
 		return "[]byte", ""
 	case FieldTypeFloatVector:
@@ -256,6 +262,9 @@ const (
 	FieldTypeString FieldType = 20
 	// FieldTypeVarChar field type varchar
 	FieldTypeVarChar FieldType = 21 // variable-length strings with a specified maximum length
+	// FieldTypeArray FieldType = 22
+	// FieldTypeJSON field type JSON
+	FieldTypeJSON FieldType = 23
 	// FieldTypeBinaryVector field type binary vector
 	FieldTypeBinaryVector FieldType = 100
 	// FieldTypeFloatVector field type float vector

go.mod

@@ -7,7 +7,7 @@ require (
 	github.com/go-faker/faker/v4 v4.1.0
 	github.com/golang/protobuf v1.5.2
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
-	github.com/milvus-io/milvus-proto/go-api v0.0.0-20230426074830-1b5e205e11ab
+	github.com/milvus-io/milvus-proto/go-api v0.0.0-20230522080721-2975bfe7a190
 	github.com/stretchr/testify v1.8.1
 	google.golang.org/grpc v1.48.0
 	google.golang.org/grpc/examples v0.0.0-20220617181431-3e7b97febc7f
@@ -29,6 +29,6 @@ require (
 	golang.org/x/sys v0.0.0-20220422013727-9388b58f7150 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	google.golang.org/genproto v0.0.0-20220503193339-ba3ae3f07e29 // indirect
-	google.golang.org/protobuf v1.28.1 // indirect
+	google.golang.org/protobuf v1.28.1
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -159,6 +159,8 @@ github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i
 github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc=
 github.com/milvus-io/milvus-proto/go-api v0.0.0-20230426074830-1b5e205e11ab h1:8XkGKAgMvftgknyMoow5DmIjHZ28A7+U5tECi+MqJmk=
 github.com/milvus-io/milvus-proto/go-api v0.0.0-20230426074830-1b5e205e11ab/go.mod h1:148qnlmZ0Fdm1Fq+Mj/OW2uDoEP25g3mjh0vMGtkgmk=
+github.com/milvus-io/milvus-proto/go-api v0.0.0-20230522080721-2975bfe7a190 h1:ZREJhOMgAvXs+K0ain51ibxAtCB8Lnn3EBZFHEykIlk=
+github.com/milvus-io/milvus-proto/go-api v0.0.0-20230522080721-2975bfe7a190/go.mod h1:148qnlmZ0Fdm1Fq+Mj/OW2uDoEP25g3mjh0vMGtkgmk=
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=