test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-11-29 18:38:44 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

enhance: [GoSDK] Add Slice method for Vector Columns (#37951)

Browse Source 
Related to #37768

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2024-11-24 14:48:33 +08:00 committed by GitHub
parent 930112e55e
commit 8c9dab5cf7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 198 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
client/column/sparse.go
View File

@ -45,3 +45,9 @@ func (c *ColumnSparseFloatVector) FieldData() *schemapb.FieldData {
vectors.Dim = int64(max.Dim())
return fd
}
func (c *ColumnSparseFloatVector) Slice(start, end int) Column {
return &ColumnSparseFloatVector{
vectorBase: c.vectorBase.slice(start, end),
}
}

9
client/column/sparse_test.go
View File

@ -91,4 +91,13 @@ func TestColumnSparseEmbedding(t *testing.T) {
assert.Equal(t, v, getV)
}
})
t.Run("test_column_slice", func(t *testing.T) {
l := rand.Intn(columnLen)
sliced := column.Slice(0, l)
slicedColumn, ok := sliced.(*ColumnSparseFloatVector)
if assert.True(t, ok) {
assert.Equal(t, column.Data()[:l], slicedColumn.Data())
}
})
}

31
client/column/vector.go
View File

@ -40,6 +40,13 @@ func (b *vectorBase[T]) FieldData() *schemapb.FieldData {
return fd
}
func (b *vectorBase[T]) slice(start, end int) *vectorBase[T] {
return &vectorBase[T]{
genericColumnBase: b.genericColumnBase.slice(start, end),
dim: b.dim,
}
}
func newVectorBase[T entity.Vector](fieldName string, dim int, vectors []T, fieldType entity.FieldType) *vectorBase[T] {
return &vectorBase[T]{
genericColumnBase: &genericColumnBase[T]{
@ -78,6 +85,12 @@ func (c *ColumnFloatVector) AppendValue(i interface{}) error {
return nil
}
func (c *ColumnFloatVector) Slice(start, end int) Column {
return &ColumnFloatVector{
vectorBase: c.vectorBase.slice(start, end),
}
}
/* binary vector */
type ColumnBinaryVector struct {
@ -105,6 +118,12 @@ func (c *ColumnBinaryVector) AppendValue(i interface{}) error {
return nil
}
func (c *ColumnBinaryVector) Slice(start, end int) Column {
return &ColumnBinaryVector{
vectorBase: c.vectorBase.slice(start, end),
}
}
/* fp16 vector */
type ColumnFloat16Vector struct {
@ -132,6 +151,12 @@ func (c *ColumnFloat16Vector) AppendValue(i interface{}) error {
return nil
}
func (c *ColumnFloat16Vector) Slice(start, end int) Column {
return &ColumnFloat16Vector{
vectorBase: c.vectorBase.slice(start, end),
}
}
type ColumnBFloat16Vector struct {
*vectorBase[entity.BFloat16Vector]
}
@ -156,3 +181,9 @@ func (c *ColumnBFloat16Vector) AppendValue(i interface{}) error {
}
return nil
}
func (c *ColumnBFloat16Vector) Slice(start, end int) Column {
return &ColumnBFloat16Vector{
vectorBase: c.vectorBase.slice(start, end),
}
}

90
client/column/vector_test.go
View File

@ -162,6 +162,96 @@ func (s *VectorSuite) TestBasic() {
})
}
func (s *VectorSuite) TestSlice() {
s.Run("float_vector", func() {
name := fmt.Sprintf("field_%d", rand.Intn(1000))
n := 100
dim := rand.Intn(10) + 2
data := make([][]float32, 0, n)
for i := 0; i < n; i++ {
row := lo.RepeatBy(dim, func(i int) float32 {
return rand.Float32()
})
data = append(data, row)
}
column := NewColumnFloatVector(name, dim, data)
l := rand.Intn(n)
sliced := column.Slice(0, l)
slicedColumn, ok := sliced.(*ColumnFloatVector)
if s.True(ok) {
s.Equal(dim, slicedColumn.Dim())
s.Equal(lo.Map(data[:l], func(row []float32, _ int) entity.FloatVector { return entity.FloatVector(row) }), slicedColumn.Data())
}
})
s.Run("binary_vector", func() {
name := fmt.Sprintf("field_%d", rand.Intn(1000))
n := 100
dim := (rand.Intn(10) + 1) * 8
data := make([][]byte, 0, n)
for i := 0; i < n; i++ {
row := lo.RepeatBy(dim/8, func(i int) byte {
return byte(rand.Intn(math.MaxUint8))
})
data = append(data, row)
}
column := NewColumnBinaryVector(name, dim, data)
l := rand.Intn(n)
sliced := column.Slice(0, l)
slicedColumn, ok := sliced.(*ColumnBinaryVector)
if s.True(ok) {
s.Equal(dim, slicedColumn.Dim())
s.Equal(lo.Map(data[:l], func(row []byte, _ int) entity.BinaryVector { return entity.BinaryVector(row) }), slicedColumn.Data())
}
})
s.Run("fp16_vector", func() {
name := fmt.Sprintf("field_%d", rand.Intn(1000))
n := 3
dim := rand.Intn(10) + 1
data := make([][]byte, 0, n)
for i := 0; i < n; i++ {
row := lo.RepeatBy(dim*2, func(i int) byte {
return byte(rand.Intn(math.MaxUint8))
})
data = append(data, row)
}
column := NewColumnFloat16Vector(name, dim, data)
l := rand.Intn(n)
sliced := column.Slice(0, l)
slicedColumn, ok := sliced.(*ColumnFloat16Vector)
if s.True(ok) {
s.Equal(dim, slicedColumn.Dim())
s.Equal(lo.Map(data[:l], func(row []byte, _ int) entity.Float16Vector { return entity.Float16Vector(row) }), slicedColumn.Data())
}
})
s.Run("bf16_vector", func() {
name := fmt.Sprintf("field_%d", rand.Intn(1000))
n := 3
dim := rand.Intn(10) + 1
data := make([][]byte, 0, n)
for i := 0; i < n; i++ {
row := lo.RepeatBy(dim*2, func(i int) byte {
return byte(rand.Intn(math.MaxUint8))
})
data = append(data, row)
}
column := NewColumnBFloat16Vector(name, dim, data)
l := rand.Intn(n)
sliced := column.Slice(0, l)
slicedColumn, ok := sliced.(*ColumnBFloat16Vector)
if s.True(ok) {
s.Equal(dim, slicedColumn.Dim())
s.Equal(lo.Map(data[:l], func(row []byte, _ int) entity.BFloat16Vector { return entity.BFloat16Vector(row) }), slicedColumn.Data())
}
})
}
func TestVectors(t *testing.T) {
suite.Run(t, new(VectorSuite))
}

2
client/go.mod
View File

@ -6,7 +6,7 @@ require (
github.com/blang/semver/v4 v4.0.0
github.com/cockroachdb/errors v1.9.1
github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20241108105827-266fb751b620
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20241120015424-93892e628c69
github.com/milvus-io/milvus/pkg v0.0.2-0.20241111021426-5e90f348fcbb
github.com/quasilyte/go-ruleguard/dsl v0.3.22
github.com/samber/lo v1.27.0

4
client/go.sum
View File

@ -318,8 +318,8 @@ github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfr
github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8=
github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc=
github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20241108105827-266fb751b620 h1:0IWUDtDloift7cQHalhdjuVkL/3qSeiXFqR7MofZBkg=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20241108105827-266fb751b620/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20241120015424-93892e628c69 h1:Qt0Bv2Fum3EX3OlkuQYHJINBzeU4oEuHy2lXSfB/gZw=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20241120015424-93892e628c69/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus/pkg v0.0.2-0.20241111021426-5e90f348fcbb h1:lMyIrG03agASB88AAwnk+NOU9V33lcBdtub/ZEv6IQU=
github.com/milvus-io/milvus/pkg v0.0.2-0.20241111021426-5e90f348fcbb/go.mod h1:w5nu1Z318AvgWQrGUYXaqLeVLu4JvCS/oYhxqctOZvU=
github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=

59
client/milvusclient/mock_milvus_server_test.go
View File

@ -4393,6 +4393,65 @@ func (_c *MilvusServiceServer_OperatePrivilegeGroup_Call) RunAndReturn(run func(
return _c
}
// OperatePrivilegeV2 provides a mock function with given fields: _a0, _a1
func (_m *MilvusServiceServer) OperatePrivilegeV2(_a0 context.Context, _a1 *milvuspb.OperatePrivilegeV2Request) (*commonpb.Status, error) {
ret := _m.Called(_a0, _a1)
if len(ret) == 0 {
panic("no return value specified for OperatePrivilegeV2")
}
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.OperatePrivilegeV2Request) (*commonpb.Status, error)); ok {
return rf(_a0, _a1)
}
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.OperatePrivilegeV2Request) *commonpb.Status); ok {
r0 = rf(_a0, _a1)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*commonpb.Status)
}
}
if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.OperatePrivilegeV2Request) error); ok {
r1 = rf(_a0, _a1)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MilvusServiceServer_OperatePrivilegeV2_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'OperatePrivilegeV2'
type MilvusServiceServer_OperatePrivilegeV2_Call struct {
*mock.Call
}
// OperatePrivilegeV2 is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *milvuspb.OperatePrivilegeV2Request
func (_e *MilvusServiceServer_Expecter) OperatePrivilegeV2(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_OperatePrivilegeV2_Call {
return &MilvusServiceServer_OperatePrivilegeV2_Call{Call: _e.mock.On("OperatePrivilegeV2", _a0, _a1)}
}
func (_c *MilvusServiceServer_OperatePrivilegeV2_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.OperatePrivilegeV2Request)) *MilvusServiceServer_OperatePrivilegeV2_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*milvuspb.OperatePrivilegeV2Request))
})
return _c
}
func (_c *MilvusServiceServer_OperatePrivilegeV2_Call) Return(_a0 *commonpb.Status, _a1 error) *MilvusServiceServer_OperatePrivilegeV2_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MilvusServiceServer_OperatePrivilegeV2_Call) RunAndReturn(run func(context.Context, *milvuspb.OperatePrivilegeV2Request) (*commonpb.Status, error)) *MilvusServiceServer_OperatePrivilegeV2_Call {
_c.Call.Return(run)
return _c
}
// OperateUserRole provides a mock function with given fields: _a0, _a1
func (_m *MilvusServiceServer) OperateUserRole(_a0 context.Context, _a1 *milvuspb.OperateUserRoleRequest) (*commonpb.Status, error) {
ret := _m.Called(_a0, _a1)