github.com/milvus-io/milvus-sdk-go/v2@v2.4.1/entity/columns_test.go (about)

     1  package entity
     2  
     3  import (
     4  	"math/rand"
     5  	"strconv"
     6  	"testing"
     7  	"time"
     8  
     9  	schema "github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
    10  	"github.com/stretchr/testify/assert"
    11  )
    12  
    13  func TestVectors(t *testing.T) {
    14  	rand.Seed(time.Now().UnixNano())
    15  	dim := rand.Intn(127) + 1
    16  
    17  	t.Run("test float vector", func(t *testing.T) {
    18  		raw := make([]float32, dim)
    19  		for i := 0; i < dim; i++ {
    20  			raw[i] = rand.Float32()
    21  		}
    22  
    23  		fv := FloatVector(raw)
    24  
    25  		assert.Equal(t, dim, fv.Dim())
    26  		assert.Equal(t, dim*4, len(fv.Serialize()))
    27  	})
    28  
    29  	t.Run("test binary vector", func(t *testing.T) {
    30  		raw := make([]byte, dim)
    31  		_, err := rand.Read(raw)
    32  		assert.Nil(t, err)
    33  
    34  		bv := BinaryVector(raw)
    35  
    36  		assert.Equal(t, dim*8, bv.Dim())
    37  		assert.ElementsMatch(t, raw, bv.Serialize())
    38  	})
    39  }
    40  
    41  func TestIDColumns(t *testing.T) {
    42  	dataLen := rand.Intn(100) + 1
    43  	base := rand.Intn(5000) // id start point
    44  
    45  	intPKCol := NewSchema().WithField(
    46  		NewField().WithName("pk").WithIsPrimaryKey(true).WithDataType(FieldTypeInt64),
    47  	)
    48  	strPKCol := NewSchema().WithField(
    49  		NewField().WithName("pk").WithIsPrimaryKey(true).WithDataType(FieldTypeVarChar),
    50  	)
    51  
    52  	t.Run("nil id", func(t *testing.T) {
    53  		col, err := IDColumns(intPKCol, nil, 0, -1)
    54  		assert.NoError(t, err)
    55  		assert.EqualValues(t, 0, col.Len())
    56  		col, err = IDColumns(strPKCol, nil, 0, -1)
    57  		assert.NoError(t, err)
    58  		assert.EqualValues(t, 0, col.Len())
    59  		idField := &schema.IDs{}
    60  		col, err = IDColumns(intPKCol, idField, 0, -1)
    61  		assert.NoError(t, err)
    62  		assert.EqualValues(t, 0, col.Len())
    63  		col, err = IDColumns(strPKCol, idField, 0, -1)
    64  		assert.NoError(t, err)
    65  		assert.EqualValues(t, 0, col.Len())
    66  	})
    67  
    68  	t.Run("int ids", func(t *testing.T) {
    69  		ids := make([]int64, 0, dataLen)
    70  		for i := 0; i < dataLen; i++ {
    71  			ids = append(ids, int64(i+base))
    72  		}
    73  		idField := &schema.IDs{
    74  			IdField: &schema.IDs_IntId{
    75  				IntId: &schema.LongArray{
    76  					Data: ids,
    77  				},
    78  			},
    79  		}
    80  		column, err := IDColumns(intPKCol, idField, 0, dataLen)
    81  		assert.Nil(t, err)
    82  		assert.NotNil(t, column)
    83  		assert.Equal(t, dataLen, column.Len())
    84  
    85  		column, err = IDColumns(intPKCol, idField, 0, -1) // test -1 method
    86  		assert.Nil(t, err)
    87  		assert.NotNil(t, column)
    88  		assert.Equal(t, dataLen, column.Len())
    89  	})
    90  	t.Run("string ids", func(t *testing.T) {
    91  		ids := make([]string, 0, dataLen)
    92  		for i := 0; i < dataLen; i++ {
    93  			ids = append(ids, strconv.FormatInt(int64(i+base), 10))
    94  		}
    95  		idField := &schema.IDs{
    96  			IdField: &schema.IDs_StrId{
    97  				StrId: &schema.StringArray{
    98  					Data: ids,
    99  				},
   100  			},
   101  		}
   102  		column, err := IDColumns(strPKCol, idField, 0, dataLen)
   103  		assert.Nil(t, err)
   104  		assert.NotNil(t, column)
   105  		assert.Equal(t, dataLen, column.Len())
   106  
   107  		column, err = IDColumns(strPKCol, idField, 0, -1) // test -1 method
   108  		assert.Nil(t, err)
   109  		assert.NotNil(t, column)
   110  		assert.Equal(t, dataLen, column.Len())
   111  	})
   112  }
   113  
   114  func TestGetIntData(t *testing.T) {
   115  	type testCase struct {
   116  		tag      string
   117  		fd       *schema.FieldData
   118  		expectOK bool
   119  	}
   120  
   121  	cases := []testCase{
   122  		{
   123  			tag: "normal_IntData",
   124  			fd: &schema.FieldData{
   125  				Field: &schema.FieldData_Scalars{
   126  					Scalars: &schema.ScalarField{
   127  						Data: &schema.ScalarField_IntData{
   128  							IntData: &schema.IntArray{Data: []int32{1, 2, 3}},
   129  						},
   130  					},
   131  				},
   132  			},
   133  			expectOK: true,
   134  		},
   135  		{
   136  			tag: "empty_LongData",
   137  			fd: &schema.FieldData{
   138  				Field: &schema.FieldData_Scalars{
   139  					Scalars: &schema.ScalarField{
   140  						Data: &schema.ScalarField_LongData{
   141  							LongData: &schema.LongArray{Data: nil},
   142  						},
   143  					},
   144  				},
   145  			},
   146  			expectOK: true,
   147  		},
   148  		{
   149  			tag: "nonempty_LongData",
   150  			fd: &schema.FieldData{
   151  				Field: &schema.FieldData_Scalars{
   152  					Scalars: &schema.ScalarField{
   153  						Data: &schema.ScalarField_LongData{
   154  							LongData: &schema.LongArray{Data: []int64{1, 2, 3}},
   155  						},
   156  					},
   157  				},
   158  			},
   159  			expectOK: false,
   160  		},
   161  		{
   162  			tag: "other_data",
   163  			fd: &schema.FieldData{
   164  				Field: &schema.FieldData_Scalars{
   165  					Scalars: &schema.ScalarField{
   166  						Data: &schema.ScalarField_BoolData{},
   167  					},
   168  				},
   169  			},
   170  			expectOK: false,
   171  		},
   172  		{
   173  			tag: "vector_data",
   174  			fd: &schema.FieldData{
   175  				Field: &schema.FieldData_Vectors{},
   176  			},
   177  			expectOK: false,
   178  		},
   179  	}
   180  
   181  	for _, tc := range cases {
   182  		t.Run(tc.tag, func(t *testing.T) {
   183  			_, ok := getIntData(tc.fd)
   184  			assert.Equal(t, tc.expectOK, ok)
   185  		})
   186  	}
   187  }