test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-11-30 02:48:45 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

feat: Bulk insert support fp16/bf16 (#32157)

Browse Source 
Issue: #22837

Signed-off-by: Cai Yudong <yudong.cai@zilliz.com>
This commit is contained in:
Cai Yudong 2024-04-22 10:05:22 +08:00 committed by GitHub
parent 037de8e4d3
commit 5fc439c600
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
15 changed files with 463 additions and 381 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
internal/storage/insert_data.go
View File

@ -616,7 +616,7 @@ func (data *BinaryVectorFieldData) AppendRows(rows interface{}) error {
// AppendRows appends FLATTEN vectors to field data.
func (data *FloatVectorFieldData) AppendRows(rows interface{}) error {
v, ok := rows.([]float32)
if !ok || len(v)%(data.Dim) != 0 {
if !ok {
return merr.WrapErrParameterInvalid("[]float32", rows, "Wrong rows type")
}
if len(v)%(data.Dim) != 0 {
@ -629,7 +629,7 @@ func (data *FloatVectorFieldData) AppendRows(rows interface{}) error {
// AppendRows appends FLATTEN vectors to field data.
func (data *Float16VectorFieldData) AppendRows(rows interface{}) error {
v, ok := rows.([]byte)
if !ok || len(v)%(data.Dim*2) != 0 {
if !ok {
return merr.WrapErrParameterInvalid("[]byte", rows, "Wrong rows type")
}
if len(v)%(data.Dim*2) != 0 {
@ -642,7 +642,7 @@ func (data *Float16VectorFieldData) AppendRows(rows interface{}) error {
// AppendRows appends FLATTEN vectors to field data.
func (data *BFloat16VectorFieldData) AppendRows(rows interface{}) error {
v, ok := rows.([]byte)
if !ok || len(v)%(data.Dim*2) != 0 {
if !ok {
return merr.WrapErrParameterInvalid("[]byte", rows, "Wrong rows type")
}
if len(v)%(data.Dim*2) != 0 {
@ -665,19 +665,17 @@ func (data *SparseFloatVectorFieldData) AppendRows(rows interface{}) error {
}
// GetMemorySize implements FieldData.GetMemorySize
func (data *BoolFieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *Int8FieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *Int16FieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *Int32FieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *Int64FieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *FloatFieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *DoubleFieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *BinaryVectorFieldData) GetMemorySize() int { return binary.Size(data.Data) + 4 }
func (data *FloatVectorFieldData) GetMemorySize() int { return binary.Size(data.Data) + 4 }
func (data *Float16VectorFieldData) GetMemorySize() int { return binary.Size(data.Data) + 4 }
func (data *BFloat16VectorFieldData) GetMemorySize() int {
return binary.Size(data.Data) + 4
}
func (data *BoolFieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *Int8FieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *Int16FieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *Int32FieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *Int64FieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *FloatFieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *DoubleFieldData) GetMemorySize() int { return binary.Size(data.Data) }
func (data *BinaryVectorFieldData) GetMemorySize() int { return binary.Size(data.Data) + 4 }
func (data *FloatVectorFieldData) GetMemorySize() int { return binary.Size(data.Data) + 4 }
func (data *Float16VectorFieldData) GetMemorySize() int { return binary.Size(data.Data) + 4 }
func (data *BFloat16VectorFieldData) GetMemorySize() int { return binary.Size(data.Data) + 4 }
func (data *SparseFloatVectorFieldData) GetMemorySize() int {
// TODO(SPARSE): should this be the memory size of serialzied size?
@ -768,9 +766,9 @@ func (data *Int64FieldData) GetRowSize(i int) int { return 8 }
func (data *FloatFieldData) GetRowSize(i int) int { return 4 }
func (data *DoubleFieldData) GetRowSize(i int) int { return 8 }
func (data *BinaryVectorFieldData) GetRowSize(i int) int { return data.Dim / 8 }
func (data *FloatVectorFieldData) GetRowSize(i int) int { return data.Dim }
func (data *Float16VectorFieldData) GetRowSize(i int) int { return data.Dim / 2 }
func (data *BFloat16VectorFieldData) GetRowSize(i int) int { return data.Dim / 2 }
func (data *FloatVectorFieldData) GetRowSize(i int) int { return data.Dim * 4 }
func (data *Float16VectorFieldData) GetRowSize(i int) int { return data.Dim * 2 }
func (data *BFloat16VectorFieldData) GetRowSize(i int) int { return data.Dim * 2 }
func (data *StringFieldData) GetRowSize(i int) int { return len(data.Data[i]) + 16 }
func (data *JSONFieldData) GetRowSize(i int) int { return len(data.Data[i]) + 16 }
func (data *ArrayFieldData) GetRowSize(i int) int {

6
internal/storage/insert_data_test.go
View File

@ -166,9 +166,9 @@ func (s *InsertDataSuite) TestGetRowSize() {
s.Equal(s.iDataOneRow.Data[JSONField].GetRowSize(0), len([]byte(`{"batch":1}`))+16)
s.Equal(s.iDataOneRow.Data[ArrayField].GetRowSize(0), 3*4)
s.Equal(s.iDataOneRow.Data[BinaryVectorField].GetRowSize(0), 1)
s.Equal(s.iDataOneRow.Data[FloatVectorField].GetRowSize(0), 4)
s.Equal(s.iDataOneRow.Data[Float16VectorField].GetRowSize(0), 2)
s.Equal(s.iDataOneRow.Data[BFloat16VectorField].GetRowSize(0), 2)
s.Equal(s.iDataOneRow.Data[FloatVectorField].GetRowSize(0), 16)
s.Equal(s.iDataOneRow.Data[Float16VectorField].GetRowSize(0), 8)
s.Equal(s.iDataOneRow.Data[BFloat16VectorField].GetRowSize(0), 8)
}
func (s *InsertDataSuite) TestGetDataType() {

18
internal/util/importutilv2/binlog/reader_test.go
View File

@ -145,6 +145,10 @@ func createBinlogBuf(t *testing.T, field *schemapb.FieldSchema, data storage.Fie
vectors := data.(*storage.Float16VectorFieldData).Data
err = evt.AddFloat16VectorToPayload(vectors, int(dim))
assert.NoError(t, err)
case schemapb.DataType_BFloat16Vector:
vectors := data.(*storage.BFloat16VectorFieldData).Data
err = evt.AddBFloat16VectorToPayload(vectors, int(dim))
assert.NoError(t, err)
default:
assert.True(t, false)
return nil
@ -242,6 +246,14 @@ func createInsertData(t *testing.T, schema *schemapb.CollectionSchema, rowCount
_, err = rand2.Read(float16VecData)
assert.NoError(t, err)
insertData.Data[field.GetFieldID()] = &storage.Float16VectorFieldData{Data: float16VecData, Dim: int(dim)}
case schemapb.DataType_BFloat16Vector:
dim, err := typeutil.GetDim(field)
assert.NoError(t, err)
total := int64(rowCount) * dim * 2
bfloat16VecData := make([]byte, total)
_, err = rand2.Read(bfloat16VecData)
assert.NoError(t, err)
insertData.Data[field.GetFieldID()] = &storage.BFloat16VectorFieldData{Data: bfloat16VecData, Dim: int(dim)}
case schemapb.DataType_String, schemapb.DataType_VarChar:
varcharData := make([]string, 0)
for i := 0; i < rowCount; i++ {
@ -441,11 +453,15 @@ func (suite *ReaderSuite) TestStringPK() {
suite.run(schemapb.DataType_Int32)
}
func (suite *ReaderSuite) TestBinaryAndFloat16Vector() {
func (suite *ReaderSuite) TestVector() {
suite.vecDataType = schemapb.DataType_BinaryVector
suite.run(schemapb.DataType_Int32)
suite.vecDataType = schemapb.DataType_FloatVector
suite.run(schemapb.DataType_Int32)
suite.vecDataType = schemapb.DataType_Float16Vector
suite.run(schemapb.DataType_Int32)
suite.vecDataType = schemapb.DataType_BFloat16Vector
suite.run(schemapb.DataType_Int32)
}
func TestUtil(t *testing.T) {

16
internal/util/importutilv2/json/reader_test.go
View File

@ -135,6 +135,14 @@ func createInsertData(t *testing.T, schema *schemapb.CollectionSchema, rowCount
_, err = rand2.Read(float16VecData)
assert.NoError(t, err)
insertData.Data[field.GetFieldID()] = &storage.Float16VectorFieldData{Data: float16VecData, Dim: int(dim)}
case schemapb.DataType_BFloat16Vector:
dim, err := typeutil.GetDim(field)
assert.NoError(t, err)
total := int64(rowCount) * dim * 2
bfloat16VecData := make([]byte, total)
_, err = rand2.Read(bfloat16VecData)
assert.NoError(t, err)
insertData.Data[field.GetFieldID()] = &storage.BFloat16VectorFieldData{Data: bfloat16VecData, Dim: int(dim)}
case schemapb.DataType_String, schemapb.DataType_VarChar:
varcharData := make([]string, 0)
for i := 0; i < rowCount; i++ {
@ -231,7 +239,7 @@ func (suite *ReaderSuite) run(dt schemapb.DataType) {
data[fieldID] = v.GetRow(i).(*schemapb.ScalarField).GetIntData().GetData()
} else if dataType == schemapb.DataType_JSON {
data[fieldID] = string(v.GetRow(i).([]byte))
} else if dataType == schemapb.DataType_BinaryVector || dataType == schemapb.DataType_Float16Vector {
} else if dataType == schemapb.DataType_BinaryVector || dataType == schemapb.DataType_Float16Vector || dataType == schemapb.DataType_BFloat16Vector {
bytes := v.GetRow(i).([]byte)
ints := make([]int, 0, len(bytes))
for _, b := range bytes {
@ -304,11 +312,15 @@ func (suite *ReaderSuite) TestStringPK() {
suite.run(schemapb.DataType_Int32)
}
func (suite *ReaderSuite) TestBinaryAndFloat16Vector() {
func (suite *ReaderSuite) TestVector() {
suite.vecDataType = schemapb.DataType_BinaryVector
suite.run(schemapb.DataType_Int32)
suite.vecDataType = schemapb.DataType_FloatVector
suite.run(schemapb.DataType_Int32)
suite.vecDataType = schemapb.DataType_Float16Vector
suite.run(schemapb.DataType_Int32)
suite.vecDataType = schemapb.DataType_BFloat16Vector
suite.run(schemapb.DataType_Int32)
}
func TestUtil(t *testing.T) {

6
internal/util/importutilv2/json/row_parser.go
View File

@ -265,7 +265,7 @@ func (r *rowParser) parseEntity(fieldID int64, obj any) (any, error) {
if !ok {
return nil, r.wrapTypeError(obj, fieldID)
}
if len(arr)*8 != r.dim {
if len(arr) != r.dim/8 {
return nil, r.wrapDimError(len(arr)*8, fieldID)
}
vec := make([]byte, len(arr))
@ -302,12 +302,12 @@ func (r *rowParser) parseEntity(fieldID int64, obj any) (any, error) {
vec[i] = float32(num)
}
return vec, nil
case schemapb.DataType_Float16Vector:
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
arr, ok := obj.([]interface{})
if !ok {
return nil, r.wrapTypeError(obj, fieldID)
}
if len(arr)/2 != r.dim {
if len(arr) != r.dim*2 {
return nil, r.wrapDimError(len(arr)/2, fieldID)
}
vec := make([]byte, len(arr))

13
internal/util/importutilv2/numpy/field_reader.go
View File

@ -94,10 +94,13 @@ func (c *FieldReader) getCount(count int64) int64 {
if total == 0 {
return 0
}
if c.field.GetDataType() == schemapb.DataType_BinaryVector {
switch c.field.GetDataType() {
case schemapb.DataType_BinaryVector:
count *= c.dim / 8
} else if c.field.GetDataType() == schemapb.DataType_FloatVector {
case schemapb.DataType_FloatVector:
count *= c.dim
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
count *= c.dim * 2
}
if int(count) > (total - c.readPosition) {
return int64(total - c.readPosition)
@ -228,6 +231,12 @@ func (c *FieldReader) Next(count int64) (any, error) {
})
}
c.readPosition += int(readCount)
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
data, err = ReadN[byte](c.reader, c.order, readCount)
if err != nil {
return nil, err
}
c.readPosition += int(readCount)
default:
return nil, merr.WrapErrImportFailed(fmt.Sprintf("unsupported data type: %s", dt.String()))
}

38
internal/util/importutilv2/numpy/reader_test.go
View File

@ -137,6 +137,14 @@ func createInsertData(t *testing.T, schema *schemapb.CollectionSchema, rowCount
_, err = rand2.Read(float16VecData)
assert.NoError(t, err)
insertData.Data[field.GetFieldID()] = &storage.Float16VectorFieldData{Data: float16VecData, Dim: int(dim)}
case schemapb.DataType_BFloat16Vector:
dim, err := typeutil.GetDim(field)
assert.NoError(t, err)
total := int64(rowCount) * dim * 2
bfloat16VecData := make([]byte, total)
_, err = rand2.Read(bfloat16VecData)
assert.NoError(t, err)
insertData.Data[field.GetFieldID()] = &storage.BFloat16VectorFieldData{Data: bfloat16VecData, Dim: int(dim)}
case schemapb.DataType_String, schemapb.DataType_VarChar:
varcharData := make([]string, 0)
for i := 0; i < rowCount; i++ {
@ -270,6 +278,17 @@ func (suite *ReaderSuite) run(dt schemapb.DataType) {
cm.EXPECT().Reader(mock.Anything, files[fieldID]).Return(&mockReader{
Reader: reader,
}, nil)
} else if dataType == schemapb.DataType_Float16Vector || dataType == schemapb.DataType_BFloat16Vector {
chunked := lo.Chunk(insertData.Data[fieldID].GetRows().([]byte), dim*2)
chunkedRows := make([][dim * 2]byte, len(chunked))
for i, innerSlice := range chunked {
copy(chunkedRows[i][:], innerSlice[:])
}
reader, err := CreateReader(chunkedRows)
suite.NoError(err)
cm.EXPECT().Reader(mock.Anything, files[fieldID]).Return(&mockReader{
Reader: reader,
}, nil)
} else if dataType == schemapb.DataType_BinaryVector {
chunked := lo.Chunk(insertData.Data[fieldID].GetRows().([]byte), dim/8)
chunkedRows := make([][dim / 8]byte, len(chunked))
@ -397,6 +416,17 @@ func (suite *ReaderSuite) failRun(dt schemapb.DataType, isDynamic bool) {
cm.EXPECT().Reader(mock.Anything, files[fieldID]).Return(&mockReader{
Reader: reader,
}, nil)
} else if dataType == schemapb.DataType_Float16Vector || dataType == schemapb.DataType_BFloat16Vector {
chunked := lo.Chunk(insertData.Data[fieldID].GetRows().([]byte), dim*2)
chunkedRows := make([][dim * 2]byte, len(chunked))
for i, innerSlice := range chunked {
copy(chunkedRows[i][:], innerSlice[:])
}
reader, err := CreateReader(chunkedRows)
suite.NoError(err)
cm.EXPECT().Reader(mock.Anything, files[fieldID]).Return(&mockReader{
Reader: reader,
}, nil)
} else if dataType == schemapb.DataType_BinaryVector {
chunked := lo.Chunk(insertData.Data[fieldID].GetRows().([]byte), dim/8)
chunkedRows := make([][dim / 8]byte, len(chunked))
@ -442,9 +472,15 @@ func (suite *ReaderSuite) TestStringPK() {
suite.run(schemapb.DataType_Int32)
}
func (suite *ReaderSuite) TestBinaryVector() {
func (suite *ReaderSuite) TestVector() {
suite.vecDataType = schemapb.DataType_BinaryVector
suite.run(schemapb.DataType_Int32)
suite.vecDataType = schemapb.DataType_FloatVector
suite.run(schemapb.DataType_Int32)
suite.vecDataType = schemapb.DataType_Float16Vector
suite.run(schemapb.DataType_Int32)
suite.vecDataType = schemapb.DataType_BFloat16Vector
suite.run(schemapb.DataType_Int32)
}
func TestUtil(t *testing.T) {

14
internal/util/importutilv2/numpy/util.go
View File

@ -205,6 +205,17 @@ func validateHeader(npyReader *npy.Reader, field *schemapb.FieldSchema, dim int)
if shape[1] != dim {
return wrapDimError(shape[1], dim, field)
}
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
// TODO: need a better way to check the element type for float16/bfloat16
if elementType != schemapb.DataType_BinaryVector {
return wrapElementTypeError(elementType, field)
}
if len(shape) != 2 {
return wrapShapeError(len(shape), 2, field)
}
if shape[1] != dim*2 {
return wrapDimError(shape[1], dim, field)
}
case schemapb.DataType_BinaryVector:
if elementType != schemapb.DataType_BinaryVector {
return wrapElementTypeError(elementType, field)
@ -219,8 +230,7 @@ func validateHeader(npyReader *npy.Reader, field *schemapb.FieldSchema, dim int)
if len(shape) != 1 {
return wrapShapeError(len(shape), 1, field)
}
case schemapb.DataType_None, schemapb.DataType_Array,
schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
case schemapb.DataType_None, schemapb.DataType_Array:
return merr.WrapErrImportFailed(fmt.Sprintf("unsupported data type: %s", field.GetDataType().String()))
default:

39
internal/util/importutilv2/parquet/field_reader.go
View File

@ -122,7 +122,7 @@ func (c *FieldReader) Next(count int64) (any, error) {
}
return byteArr, nil
case schemapb.DataType_BinaryVector:
return ReadBinaryData(c, count)
return ReadBinaryData(c, schemapb.DataType_BinaryVector, count)
case schemapb.DataType_FloatVector:
arrayData, err := ReadIntegerOrFloatArrayData[float32](c, count)
if err != nil {
@ -133,6 +133,10 @@ func (c *FieldReader) Next(count int64) (any, error) {
}
vectors := lo.Flatten(arrayData.([][]float32))
return vectors, nil
case schemapb.DataType_Float16Vector:
return ReadBinaryData(c, schemapb.DataType_Float16Vector, count)
case schemapb.DataType_BFloat16Vector:
return ReadBinaryData(c, schemapb.DataType_BFloat16Vector, count)
case schemapb.DataType_Array:
data := make([]*schemapb.ScalarField, 0, count)
elementType := c.field.GetElementType()
@ -154,7 +158,6 @@ func (c *FieldReader) Next(count int64) (any, error) {
},
})
}
case schemapb.DataType_Int8:
int8Array, err := ReadIntegerOrFloatArrayData[int32](c, count)
if err != nil {
@ -172,7 +175,6 @@ func (c *FieldReader) Next(count int64) (any, error) {
},
})
}
case schemapb.DataType_Int16:
int16Array, err := ReadIntegerOrFloatArrayData[int32](c, count)
if err != nil {
@ -190,7 +192,6 @@ func (c *FieldReader) Next(count int64) (any, error) {
},
})
}
case schemapb.DataType_Int32:
int32Array, err := ReadIntegerOrFloatArrayData[int32](c, count)
if err != nil {
@ -208,7 +209,6 @@ func (c *FieldReader) Next(count int64) (any, error) {
},
})
}
case schemapb.DataType_Int64:
int64Array, err := ReadIntegerOrFloatArrayData[int64](c, count)
if err != nil {
@ -226,7 +226,6 @@ func (c *FieldReader) Next(count int64) (any, error) {
},
})
}
case schemapb.DataType_Float:
float32Array, err := ReadIntegerOrFloatArrayData[float32](c, count)
if err != nil {
@ -244,7 +243,6 @@ func (c *FieldReader) Next(count int64) (any, error) {
},
})
}
case schemapb.DataType_Double:
float64Array, err := ReadIntegerOrFloatArrayData[float64](c, count)
if err != nil {
@ -262,7 +260,6 @@ func (c *FieldReader) Next(count int64) (any, error) {
},
})
}
case schemapb.DataType_VarChar, schemapb.DataType_String:
stringArray, err := ReadStringArrayData(c, count)
if err != nil {
@ -392,7 +389,7 @@ func ReadStringData(pcr *FieldReader, count int64) (any, error) {
return data, nil
}
func ReadBinaryData(pcr *FieldReader, count int64) (any, error) {
func ReadBinaryData(pcr *FieldReader, dataType schemapb.DataType, count int64) (any, error) {
chunked, err := pcr.columnReader.NextBatch(count)
if err != nil {
return nil, err
@ -408,8 +405,8 @@ func ReadBinaryData(pcr *FieldReader, count int64) (any, error) {
}
case arrow.LIST:
listReader := chunk.(*array.List)
if !isRegularVector(listReader.Offsets(), pcr.dim, true) {
return nil, merr.WrapErrImportFailed("binary vector is irregular")
if !isVectorAligned(listReader.Offsets(), pcr.dim, dataType) {
return nil, merr.WrapErrImportFailed("%s not aligned", dataType.String())
}
uint8Reader, ok := listReader.ListValues().(*array.Uint8)
if !ok {
@ -428,15 +425,21 @@ func ReadBinaryData(pcr *FieldReader, count int64) (any, error) {
return data, nil
}
func isRegularVector(offsets []int32, dim int, isBinary bool) bool {
func isVectorAligned(offsets []int32, dim int, dataType schemapb.DataType) bool {
if len(offsets) < 1 {
return false
}
if isBinary {
dim = dim / 8
var elemCount int = 0
switch dataType {
case schemapb.DataType_BinaryVector:
elemCount = dim / 8
case schemapb.DataType_FloatVector:
elemCount = dim
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
elemCount = dim * 2
}
for i := 1; i < len(offsets); i++ {
if offsets[i]-offsets[i-1] != int32(dim) {
if offsets[i]-offsets[i-1] != int32(elemCount) {
return false
}
}
@ -497,9 +500,9 @@ func ReadIntegerOrFloatArrayData[T constraints.Integer | constraints.Float](pcr
return nil, WrapTypeErr("list", chunk.DataType().Name(), pcr.field)
}
offsets := listReader.Offsets()
if typeutil.IsVectorType(pcr.field.GetDataType()) &&
!isRegularVector(offsets, pcr.dim, pcr.field.GetDataType() == schemapb.DataType_BinaryVector) {
return nil, merr.WrapErrImportFailed("float vector is irregular")
dataType := pcr.field.GetDataType()
if typeutil.IsVectorType(dataType) && !isVectorAligned(offsets, pcr.dim, dataType) {
return nil, merr.WrapErrImportFailed("%s not aligned", dataType.String())
}
valueReader := listReader.ListValues()
switch valueReader.DataType().ID() {

132
internal/util/importutilv2/parquet/reader_test.go
View File

@ -62,84 +62,6 @@ func (s *ReaderSuite) SetupTest() {
s.vecDataType = schemapb.DataType_FloatVector
}
func milvusDataTypeToArrowType(dataType schemapb.DataType, isBinary bool) arrow.DataType {
switch dataType {
case schemapb.DataType_Bool:
return &arrow.BooleanType{}
case schemapb.DataType_Int8:
return &arrow.Int8Type{}
case schemapb.DataType_Int16:
return &arrow.Int16Type{}
case schemapb.DataType_Int32:
return &arrow.Int32Type{}
case schemapb.DataType_Int64:
return &arrow.Int64Type{}
case schemapb.DataType_Float:
return &arrow.Float32Type{}
case schemapb.DataType_Double:
return &arrow.Float64Type{}
case schemapb.DataType_VarChar, schemapb.DataType_String:
return &arrow.StringType{}
case schemapb.DataType_Array:
return &arrow.ListType{}
case schemapb.DataType_JSON:
return &arrow.StringType{}
case schemapb.DataType_FloatVector:
return arrow.ListOfField(arrow.Field{
Name: "item",
Type: &arrow.Float32Type{},
Nullable: true,
Metadata: arrow.Metadata{},
})
case schemapb.DataType_BinaryVector:
if isBinary {
return &arrow.BinaryType{}
}
return arrow.ListOfField(arrow.Field{
Name: "item",
Type: &arrow.Uint8Type{},
Nullable: true,
Metadata: arrow.Metadata{},
})
case schemapb.DataType_Float16Vector:
return arrow.ListOfField(arrow.Field{
Name: "item",
Type: &arrow.Float16Type{},
Nullable: true,
Metadata: arrow.Metadata{},
})
default:
panic("unsupported data type")
}
}
func convertMilvusSchemaToArrowSchema(schema *schemapb.CollectionSchema) *arrow.Schema {
fields := make([]arrow.Field, 0)
for _, field := range schema.GetFields() {
if field.GetDataType() == schemapb.DataType_Array {
fields = append(fields, arrow.Field{
Name: field.GetName(),
Type: arrow.ListOfField(arrow.Field{
Name: "item",
Type: milvusDataTypeToArrowType(field.GetElementType(), false),
Nullable: true,
Metadata: arrow.Metadata{},
}),
Nullable: true,
Metadata: arrow.Metadata{},
})
continue
}
fields = append(fields, arrow.Field{
Name: field.GetName(),
Type: milvusDataTypeToArrowType(field.GetDataType(), field.Name == "FieldBinaryVector2"),
Nullable: true,
Metadata: arrow.Metadata{},
})
}
return arrow.NewSchema(fields, nil)
}
func randomString(length int) string {
letterRunes := []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
b := make([]rune, length)
@ -257,6 +179,40 @@ func buildArrayData(schema *schemapb.CollectionSchema, rows int) ([]arrow.Array,
insertData.Data[field.GetFieldID()] = &storage.FloatVectorFieldData{Data: floatVecData, Dim: dim}
builder.AppendValues(offsets, valid)
columns = append(columns, builder.NewListArray())
case schemapb.DataType_Float16Vector:
float16VecData := make([]byte, 0)
builder := array.NewListBuilder(mem, &arrow.Uint8Type{})
offsets := make([]int32, 0, rows)
valid := make([]bool, 0, rows)
rowBytes := dim * 2
for i := 0; i < rowBytes*rows; i++ {
float16VecData = append(float16VecData, uint8(i%256))
}
builder.ValueBuilder().(*array.Uint8Builder).AppendValues(float16VecData, nil)
for i := 0; i < rows; i++ {
offsets = append(offsets, int32(rowBytes*i))
valid = append(valid, true)
}
insertData.Data[field.GetFieldID()] = &storage.Float16VectorFieldData{Data: float16VecData, Dim: dim}
builder.AppendValues(offsets, valid)
columns = append(columns, builder.NewListArray())
case schemapb.DataType_BFloat16Vector:
bfloat16VecData := make([]byte, 0)
builder := array.NewListBuilder(mem, &arrow.Uint8Type{})
offsets := make([]int32, 0, rows)
valid := make([]bool, 0, rows)
rowBytes := dim * 2
for i := 0; i < rowBytes*rows; i++ {
bfloat16VecData = append(bfloat16VecData, uint8(i%256))
}
builder.ValueBuilder().(*array.Uint8Builder).AppendValues(bfloat16VecData, nil)
for i := 0; i < rows; i++ {
offsets = append(offsets, int32(rowBytes*i))
valid = append(valid, true)
}
insertData.Data[field.GetFieldID()] = &storage.BFloat16VectorFieldData{Data: bfloat16VecData, Dim: dim}
builder.AppendValues(offsets, valid)
columns = append(columns, builder.NewListArray())
case schemapb.DataType_BinaryVector:
if isBinary {
binVecData := make([][]byte, 0)
@ -276,12 +232,13 @@ func buildArrayData(schema *schemapb.CollectionSchema, rows int) ([]arrow.Array,
builder := array.NewListBuilder(mem, &arrow.Uint8Type{})
offsets := make([]int32, 0, rows)
valid := make([]bool, 0)
for i := 0; i < dim*rows/8; i++ {
rowBytes := dim / 8
for i := 0; i < rowBytes*rows; i++ {
binVecData = append(binVecData, uint8(i))
}
builder.ValueBuilder().(*array.Uint8Builder).AppendValues(binVecData, nil)
for i := 0; i < rows; i++ {
offsets = append(offsets, int32(dim*i/8))
offsets = append(offsets, int32(rowBytes*i))
valid = append(valid, true)
}
builder.AppendValues(offsets, valid)
@ -472,7 +429,10 @@ func buildArrayData(schema *schemapb.CollectionSchema, rows int) ([]arrow.Array,
}
func writeParquet(w io.Writer, schema *schemapb.CollectionSchema, numRows int) (*storage.InsertData, error) {
pqSchema := convertMilvusSchemaToArrowSchema(schema)
pqSchema, err := ConvertToArrowSchema(schema)
if err != nil {
return nil, err
}
fw, err := pqarrow.NewFileWriter(pqSchema, w, parquet.NewWriterProperties(parquet.WithMaxRowGroupLength(int64(numRows))), pqarrow.DefaultWriterProps())
if err != nil {
return nil, err
@ -648,11 +608,15 @@ func (s *ReaderSuite) TestStringPK() {
s.run(schemapb.DataType_Int32)
}
func (s *ReaderSuite) TestBinaryAndFloat16Vector() {
func (s *ReaderSuite) TestVector() {
s.vecDataType = schemapb.DataType_BinaryVector
s.run(schemapb.DataType_Int32)
// s.vecDataType = schemapb.DataType_Float16Vector
// s.run(schemapb.DataType_Int32) // TODO: dyh, support float16 vector
s.vecDataType = schemapb.DataType_FloatVector
s.run(schemapb.DataType_Int32)
s.vecDataType = schemapb.DataType_Float16Vector
s.run(schemapb.DataType_Int32)
s.vecDataType = schemapb.DataType_BFloat16Vector
s.run(schemapb.DataType_Int32)
}
func TestUtil(t *testing.T) {

249
internal/util/importutilv2/parquet/util.go
View File

@ -52,6 +52,11 @@ func CreateFieldReaders(ctx context.Context, fileReader *pqarrow.FileReader, sch
return nil, merr.WrapErrImportFailed(fmt.Sprintf("get parquet schema failed, err=%v", err))
}
err = isSchemaEqual(schema, pqSchema)
if err != nil {
return nil, merr.WrapErrImportFailed(fmt.Sprintf("schema not equal, err=%v", err))
}
crs := make(map[int64]*FieldReader)
for i, pqField := range pqSchema.Fields() {
field, ok := nameToField[pqField.Name]
@ -60,28 +65,11 @@ func CreateFieldReaders(ctx context.Context, fileReader *pqarrow.FileReader, sch
return nil, merr.WrapErrImportFailed(fmt.Sprintf("the field: %s is not in schema, "+
"if it's a dynamic field, please reformat data by bulk_writer", pqField.Name))
}
if field.GetIsPrimaryKey() && field.GetAutoID() {
if typeutil.IsAutoPKField(field) {
return nil, merr.WrapErrImportFailed(
fmt.Sprintf("the primary key '%s' is auto-generated, no need to provide", field.GetName()))
}
arrowType, isList := convertArrowSchemaToDataType(pqField, false)
dataType := field.GetDataType()
if isList {
if !typeutil.IsVectorType(dataType) && dataType != schemapb.DataType_Array {
return nil, WrapTypeErr(dataType.String(), pqField.Type.Name(), field)
}
if dataType == schemapb.DataType_Array {
dataType = field.GetElementType()
}
}
if !isConvertible(arrowType, dataType, isList) {
if isList {
return nil, WrapTypeErr(dataType.String(), pqField.Type.(*arrow.ListType).ElemField().Type.Name(), field)
}
return nil, WrapTypeErr(dataType.String(), pqField.Type.Name(), field)
}
cr, err := NewFieldReader(ctx, fileReader, i, field)
if err != nil {
return nil, err
@ -94,7 +82,7 @@ func CreateFieldReaders(ctx context.Context, fileReader *pqarrow.FileReader, sch
}
for _, field := range nameToField {
if (field.GetIsPrimaryKey() && field.GetAutoID()) || field.GetIsDynamic() {
if typeutil.IsAutoPKField(field) || field.GetIsDynamic() {
continue
}
if _, ok := crs[field.GetFieldID()]; !ok {
@ -105,79 +93,166 @@ func CreateFieldReaders(ctx context.Context, fileReader *pqarrow.FileReader, sch
return crs, nil
}
func convertArrowSchemaToDataType(field arrow.Field, isList bool) (schemapb.DataType, bool) {
switch field.Type.ID() {
case arrow.BOOL:
return schemapb.DataType_Bool, false
case arrow.UINT8:
if isList {
return schemapb.DataType_BinaryVector, false
}
return schemapb.DataType_None, false
case arrow.INT8:
return schemapb.DataType_Int8, false
case arrow.INT16:
return schemapb.DataType_Int16, false
case arrow.INT32:
return schemapb.DataType_Int32, false
case arrow.INT64:
return schemapb.DataType_Int64, false
case arrow.FLOAT16:
if isList {
return schemapb.DataType_Float16Vector, false
}
return schemapb.DataType_None, false
case arrow.FLOAT32:
return schemapb.DataType_Float, false
case arrow.FLOAT64:
return schemapb.DataType_Double, false
case arrow.STRING:
return schemapb.DataType_VarChar, false
case arrow.BINARY:
return schemapb.DataType_BinaryVector, false
case arrow.LIST:
elementType, _ := convertArrowSchemaToDataType(field.Type.(*arrow.ListType).ElemField(), true)
return elementType, true
default:
return schemapb.DataType_None, false
}
}
func isConvertible(src, dst schemapb.DataType, isList bool) bool {
switch src {
case schemapb.DataType_Bool:
return typeutil.IsBoolType(dst)
case schemapb.DataType_Int8:
return typeutil.IsArithmetic(dst)
case schemapb.DataType_Int16:
return typeutil.IsArithmetic(dst) && dst != schemapb.DataType_Int8
case schemapb.DataType_Int32:
return typeutil.IsArithmetic(dst) && dst != schemapb.DataType_Int8 && dst != schemapb.DataType_Int16
case schemapb.DataType_Int64:
return typeutil.IsFloatingType(dst) || dst == schemapb.DataType_Int64
case schemapb.DataType_Float:
if isList && dst == schemapb.DataType_FloatVector {
return true
}
return typeutil.IsFloatingType(dst)
case schemapb.DataType_Double:
if isList && dst == schemapb.DataType_FloatVector {
return true
}
return dst == schemapb.DataType_Double
case schemapb.DataType_String, schemapb.DataType_VarChar:
return typeutil.IsStringType(dst) || typeutil.IsJSONType(dst)
case schemapb.DataType_JSON:
return typeutil.IsJSONType(dst)
case schemapb.DataType_BinaryVector:
return dst == schemapb.DataType_BinaryVector
case schemapb.DataType_Float16Vector:
return dst == schemapb.DataType_Float16Vector
func isArrowIntegerType(dataType arrow.Type) bool {
switch dataType {
case arrow.INT8, arrow.INT16, arrow.INT32, arrow.INT64:
return true
default:
return false
}
}
func isArrowFloatingType(dataType arrow.Type) bool {
switch dataType {
case arrow.FLOAT32, arrow.FLOAT64:
return true
default:
return false
}
}
func isArrowArithmeticType(dataType arrow.Type) bool {
return isArrowIntegerType(dataType) || isArrowFloatingType(dataType)
}
func isArrowDataTypeConvertible(src arrow.DataType, dst arrow.DataType) bool {
srcType := src.ID()
dstType := dst.ID()
switch srcType {
case arrow.BOOL:
return dstType == arrow.BOOL
case arrow.UINT8:
return dstType == arrow.UINT8
case arrow.INT8:
return isArrowArithmeticType(dstType)
case arrow.INT16:
return isArrowArithmeticType(dstType) && dstType != arrow.INT8
case arrow.INT32:
return isArrowArithmeticType(dstType) && dstType != arrow.INT8 && dstType != arrow.INT16
case arrow.INT64:
return isArrowFloatingType(dstType) || dstType == arrow.INT64
case arrow.FLOAT32:
return isArrowFloatingType(dstType)
case arrow.FLOAT64:
// TODO caiyd: need do strict type check
// return dstType == arrow.FLOAT64
return isArrowFloatingType(dstType)
case arrow.STRING:
return dstType == arrow.STRING
case arrow.BINARY:
return dstType == arrow.LIST && dst.(*arrow.ListType).Elem().ID() == arrow.UINT8
case arrow.LIST:
return dstType == arrow.LIST && isArrowDataTypeConvertible(src.(*arrow.ListType).Elem(), dst.(*arrow.ListType).Elem())
default:
return false
}
}
func convertToArrowDataType(field *schemapb.FieldSchema, isArray bool) (arrow.DataType, error) {
dataType := field.GetDataType()
if isArray {
dataType = field.GetElementType()
}
switch dataType {
case schemapb.DataType_Bool:
return &arrow.BooleanType{}, nil
case schemapb.DataType_Int8:
return &arrow.Int8Type{}, nil
case schemapb.DataType_Int16:
return &arrow.Int16Type{}, nil
case schemapb.DataType_Int32:
return &arrow.Int32Type{}, nil
case schemapb.DataType_Int64:
return &arrow.Int64Type{}, nil
case schemapb.DataType_Float:
return &arrow.Float32Type{}, nil
case schemapb.DataType_Double:
return &arrow.Float64Type{}, nil
case schemapb.DataType_VarChar, schemapb.DataType_String:
return &arrow.StringType{}, nil
case schemapb.DataType_JSON:
return &arrow.StringType{}, nil
case schemapb.DataType_Array:
elemType, err := convertToArrowDataType(field, true)
if err != nil {
return nil, err
}
return arrow.ListOfField(arrow.Field{
Name: "item",
Type: elemType,
Nullable: true,
Metadata: arrow.Metadata{},
}), nil
case schemapb.DataType_BinaryVector:
return arrow.ListOfField(arrow.Field{
Name: "item",
Type: &arrow.Uint8Type{},
Nullable: true,
Metadata: arrow.Metadata{},
}), nil
case schemapb.DataType_FloatVector:
return arrow.ListOfField(arrow.Field{
Name: "item",
Type: &arrow.Float32Type{},
Nullable: true,
Metadata: arrow.Metadata{},
}), nil
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
return arrow.ListOfField(arrow.Field{
Name: "item",
Type: &arrow.Uint8Type{},
Nullable: true,
Metadata: arrow.Metadata{},
}), nil
default:
return nil, merr.WrapErrParameterInvalidMsg("unsupported data type %v", dataType.String())
}
}
func ConvertToArrowSchema(schema *schemapb.CollectionSchema) (*arrow.Schema, error) {
arrFields := make([]arrow.Field, 0)
for _, field := range schema.GetFields() {
if typeutil.IsAutoPKField(field) {
continue
}
arrDataType, err := convertToArrowDataType(field, false)
if err != nil {
return nil, err
}
arrFields = append(arrFields, arrow.Field{
Name: field.GetName(),
Type: arrDataType,
Nullable: true,
Metadata: arrow.Metadata{},
})
}
return arrow.NewSchema(arrFields, nil), nil
}
func isSchemaEqual(schema *schemapb.CollectionSchema, arrSchema *arrow.Schema) error {
arrNameToField := lo.KeyBy(arrSchema.Fields(), func(field arrow.Field) string {
return field.Name
})
for _, field := range schema.GetFields() {
if typeutil.IsAutoPKField(field) {
continue
}
arrField, ok := arrNameToField[field.GetName()]
if !ok {
return merr.WrapErrImportFailed(fmt.Sprintf("field '%s' not in arrow schema", field.GetName()))
}
toArrDataType, err := convertToArrowDataType(field, false)
if err != nil {
return err
}
if !isArrowDataTypeConvertible(arrField.Type, toArrDataType) {
return merr.WrapErrImportFailed(fmt.Sprintf("field '%s' type mis-match, milvus data type '%s', arrow data type get '%s'",
field.Name, field.DataType.String(), arrField.Type.String()))
}
}
return nil
}
func estimateReadCountPerBatch(bufferSize int, schema *schemapb.CollectionSchema) (int64, error) {
sizePerRecord, err := typeutil.EstimateMaxSizePerRecord(schema)
if err != nil {

17
pkg/util/typeutil/schema.go
View File

@ -137,18 +137,7 @@ func estimateSizeBy(schema *schemapb.CollectionSchema, policy getVariableFieldLe
break
}
}
case schemapb.DataType_Float16Vector:
for _, kv := range fs.TypeParams {
if kv.Key == common.DimKey {
v, err := strconv.Atoi(kv.Value)
if err != nil {
return -1, err
}
res += v * 2
break
}
}
case schemapb.DataType_BFloat16Vector:
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
for _, kv := range fs.TypeParams {
if kv.Key == common.DimKey {
v, err := strconv.Atoi(kv.Value)
@ -1138,6 +1127,10 @@ func IsPrimaryFieldDataExist(datas []*schemapb.FieldData, primaryFieldSchema *sc
return primaryFieldData != nil
}
func IsAutoPKField(field *schemapb.FieldSchema) bool {
return field.GetIsPrimaryKey() && field.GetAutoID()
}
func AppendSystemFields(schema *schemapb.CollectionSchema) *schemapb.CollectionSchema {
newSchema := proto.Clone(schema).(*schemapb.CollectionSchema)
newSchema.Fields = append(newSchema.Fields, &schemapb.FieldSchema{

60
tests/integration/import/import_test.go
View File

@ -36,6 +36,7 @@ import (
"github.com/milvus-io/milvus/pkg/common"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util/funcutil"
"github.com/milvus-io/milvus/pkg/util/indexparamcheck"
"github.com/milvus-io/milvus/pkg/util/metric"
"github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/milvus-io/milvus/tests/integration"
@ -50,6 +51,10 @@ type BulkInsertSuite struct {
pkType schemapb.DataType
autoID bool
fileType importutilv2.FileType
vecType schemapb.DataType
indexType indexparamcheck.IndexType
metricType metric.MetricType
}
func (s *BulkInsertSuite) SetupTest() {
@ -59,6 +64,10 @@ func (s *BulkInsertSuite) SetupTest() {
s.fileType = importutilv2.Parquet
s.pkType = schemapb.DataType_Int64
s.autoID = false
s.vecType = schemapb.DataType_FloatVector
s.indexType = indexparamcheck.IndexHNSW
s.metricType = metric.L2
}
func (s *BulkInsertSuite) run() {
@ -75,7 +84,7 @@ func (s *BulkInsertSuite) run() {
schema := integration.ConstructSchema(collectionName, dim, s.autoID,
&schemapb.FieldSchema{FieldID: 100, Name: "id", DataType: s.pkType, TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxLengthKey, Value: "128"}}, IsPrimaryKey: true, AutoID: s.autoID},
&schemapb.FieldSchema{FieldID: 101, Name: "image_path", DataType: schemapb.DataType_VarChar, TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxLengthKey, Value: "65535"}}},
&schemapb.FieldSchema{FieldID: 102, Name: "embeddings", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "128"}}},
&schemapb.FieldSchema{FieldID: 102, Name: "embeddings", DataType: s.vecType, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "128"}}},
)
marshaledSchema, err := proto.Marshal(schema)
s.NoError(err)
@ -91,11 +100,13 @@ func (s *BulkInsertSuite) run() {
var files []*internalpb.ImportFile
err = os.MkdirAll(c.ChunkManager.RootPath(), os.ModePerm)
s.NoError(err)
if s.fileType == importutilv2.Numpy {
switch s.fileType {
case importutilv2.Numpy:
importFile, err := GenerateNumpyFiles(c.ChunkManager, schema, rowCount)
s.NoError(err)
files = []*internalpb.ImportFile{importFile}
} else if s.fileType == importutilv2.JSON {
case importutilv2.JSON:
rowBasedFile := c.ChunkManager.RootPath() + "/" + "test.json"
GenerateJSONFile(s.T(), rowBasedFile, schema, rowCount)
defer os.Remove(rowBasedFile)
@ -106,7 +117,7 @@ func (s *BulkInsertSuite) run() {
},
},
}
} else if s.fileType == importutilv2.Parquet {
case importutilv2.Parquet:
filePath := fmt.Sprintf("/tmp/test_%d.parquet", rand.Int())
err = GenerateParquetFile(filePath, schema, rowCount)
s.NoError(err)
@ -147,7 +158,7 @@ func (s *BulkInsertSuite) run() {
CollectionName: collectionName,
FieldName: "embeddings",
IndexName: "_default",
ExtraParams: integration.ConstructIndexParam(dim, integration.IndexHNSW, metric.L2),
ExtraParams: integration.ConstructIndexParam(dim, s.indexType, s.metricType),
})
s.NoError(err)
s.Equal(commonpb.ErrorCode_Success, createIndexStatus.GetErrorCode())
@ -168,28 +179,41 @@ func (s *BulkInsertSuite) run() {
topk := 10
roundDecimal := -1
params := integration.GetSearchParams(integration.IndexHNSW, metric.L2)
params := integration.GetSearchParams(s.indexType, s.metricType)
searchReq := integration.ConstructSearchRequest("", collectionName, expr,
"embeddings", schemapb.DataType_FloatVector, nil, metric.L2, params, nq, dim, topk, roundDecimal)
"embeddings", s.vecType, nil, s.metricType, params, nq, dim, topk, roundDecimal)
searchResult, err := c.Proxy.Search(ctx, searchReq)
s.NoError(err)
s.Equal(commonpb.ErrorCode_Success, searchResult.GetStatus().GetErrorCode())
}
func (s *BulkInsertSuite) TestNumpy() {
s.fileType = importutilv2.Numpy
s.run()
}
func (s *BulkInsertSuite) TestMultiFileTypes() {
fileTypeArr := []importutilv2.FileType{importutilv2.JSON, importutilv2.Numpy, importutilv2.Parquet}
func (s *BulkInsertSuite) TestJSON() {
s.fileType = importutilv2.JSON
s.run()
}
for _, fileType := range fileTypeArr {
s.fileType = fileType
func (s *BulkInsertSuite) TestParquet() {
s.fileType = importutilv2.Parquet
s.run()
s.vecType = schemapb.DataType_BinaryVector
s.indexType = indexparamcheck.IndexFaissBinIvfFlat
s.metricType = metric.HAMMING
s.run()
s.vecType = schemapb.DataType_FloatVector
s.indexType = indexparamcheck.IndexHNSW
s.metricType = metric.L2
s.run()
s.vecType = schemapb.DataType_Float16Vector
s.indexType = indexparamcheck.IndexHNSW
s.metricType = metric.L2
s.run()
s.vecType = schemapb.DataType_BFloat16Vector
s.indexType = indexparamcheck.IndexHNSW
s.metricType = metric.L2
s.run()
}
}
func (s *BulkInsertSuite) TestAutoID() {

177
tests/integration/import/util_test.go
View File

@ -40,6 +40,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/storage"
pq "github.com/milvus-io/milvus/internal/util/importutilv2/parquet"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/typeutil"
@ -124,6 +125,14 @@ func createInsertData(t *testing.T, schema *schemapb.CollectionSchema, rowCount
_, err = rand2.Read(float16VecData)
assert.NoError(t, err)
insertData.Data[field.GetFieldID()] = &storage.Float16VectorFieldData{Data: float16VecData, Dim: int(dim)}
case schemapb.DataType_BFloat16Vector:
dim, err := typeutil.GetDim(field)
assert.NoError(t, err)
total := int64(rowCount) * dim * 2
bfloat16VecData := make([]byte, total)
_, err = rand2.Read(bfloat16VecData)
assert.NoError(t, err)
insertData.Data[field.GetFieldID()] = &storage.BFloat16VectorFieldData{Data: bfloat16VecData, Dim: int(dim)}
case schemapb.DataType_String, schemapb.DataType_VarChar:
varcharData := make([]string, 0)
for i := 0; i < rowCount; i++ {
@ -155,87 +164,6 @@ func createInsertData(t *testing.T, schema *schemapb.CollectionSchema, rowCount
return insertData
}
func milvusDataTypeToArrowType(dataType schemapb.DataType, isBinary bool) arrow.DataType {
switch dataType {
case schemapb.DataType_Bool:
return &arrow.BooleanType{}
case schemapb.DataType_Int8:
return &arrow.Int8Type{}
case schemapb.DataType_Int16:
return &arrow.Int16Type{}
case schemapb.DataType_Int32:
return &arrow.Int32Type{}
case schemapb.DataType_Int64:
return &arrow.Int64Type{}
case schemapb.DataType_Float:
return &arrow.Float32Type{}
case schemapb.DataType_Double:
return &arrow.Float64Type{}
case schemapb.DataType_VarChar, schemapb.DataType_String:
return &arrow.StringType{}
case schemapb.DataType_Array:
return &arrow.ListType{}
case schemapb.DataType_JSON:
return &arrow.StringType{}
case schemapb.DataType_FloatVector:
return arrow.ListOfField(arrow.Field{
Name: "item",
Type: &arrow.Float32Type{},
Nullable: true,
Metadata: arrow.Metadata{},
})
case schemapb.DataType_BinaryVector:
if isBinary {
return &arrow.BinaryType{}
}
return arrow.ListOfField(arrow.Field{
Name: "item",
Type: &arrow.Uint8Type{},
Nullable: true,
Metadata: arrow.Metadata{},
})
case schemapb.DataType_Float16Vector:
return arrow.ListOfField(arrow.Field{
Name: "item",
Type: &arrow.Float16Type{},
Nullable: true,
Metadata: arrow.Metadata{},
})
default:
panic("unsupported data type")
}
}
func convertMilvusSchemaToArrowSchema(schema *schemapb.CollectionSchema) *arrow.Schema {
fields := make([]arrow.Field, 0)
for _, field := range schema.GetFields() {
if field.GetIsPrimaryKey() && field.GetAutoID() {
continue
}
if field.GetDataType() == schemapb.DataType_Array {
fields = append(fields, arrow.Field{
Name: field.GetName(),
Type: arrow.ListOfField(arrow.Field{
Name: "item",
Type: milvusDataTypeToArrowType(field.GetElementType(), false),
Nullable: true,
Metadata: arrow.Metadata{},
}),
Nullable: true,
Metadata: arrow.Metadata{},
})
continue
}
fields = append(fields, arrow.Field{
Name: field.GetName(),
Type: milvusDataTypeToArrowType(field.GetDataType(), field.Name == "FieldBinaryVector2"),
Nullable: true,
Metadata: arrow.Metadata{},
})
}
return arrow.NewSchema(fields, nil)
}
func randomString(length int) string {
letterRunes := []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
b := make([]rune, length)
@ -245,7 +173,7 @@ func randomString(length int) string {
return string(b)
}
func buildArrayData(dataType, elementType schemapb.DataType, dim, rows int, isBinary bool) arrow.Array {
func buildArrayData(dataType, elemType schemapb.DataType, dim, rows int) arrow.Array {
mem := memory.NewGoAllocator()
switch dataType {
case schemapb.DataType_Bool:
@ -296,6 +224,20 @@ func buildArrayData(dataType, elementType schemapb.DataType, dim, rows int, isBi
builder.Append(randomString(10))
}
return builder.NewStringArray()
case schemapb.DataType_BinaryVector:
builder := array.NewListBuilder(mem, &arrow.Uint8Type{})
offsets := make([]int32, 0, rows)
valid := make([]bool, 0)
rowBytes := dim / 8
for i := 0; i < rowBytes*rows; i++ {
builder.ValueBuilder().(*array.Uint8Builder).Append(uint8(i % 256))
}
for i := 0; i < rows; i++ {
offsets = append(offsets, int32(rowBytes*i))
valid = append(valid, true)
}
builder.AppendValues(offsets, valid)
return builder.NewListArray()
case schemapb.DataType_FloatVector:
builder := array.NewListBuilder(mem, &arrow.Float32Type{})
offsets := make([]int32, 0, rows)
@ -304,31 +246,21 @@ func buildArrayData(dataType, elementType schemapb.DataType, dim, rows int, isBi
builder.ValueBuilder().(*array.Float32Builder).Append(float32(i))
}
for i := 0; i < rows; i++ {
offsets = append(offsets, int32(i*dim))
offsets = append(offsets, int32(dim*i))
valid = append(valid, true)
}
builder.AppendValues(offsets, valid)
return builder.NewListArray()
case schemapb.DataType_BinaryVector:
if isBinary {
builder := array.NewBinaryBuilder(mem, &arrow.BinaryType{})
for i := 0; i < rows; i++ {
element := make([]byte, dim/8)
for j := 0; j < dim/8; j++ {
element[j] = randomString(1)[0]
}
builder.Append(element)
}
return builder.NewBinaryArray()
}
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
builder := array.NewListBuilder(mem, &arrow.Uint8Type{})
offsets := make([]int32, 0, rows)
valid := make([]bool, 0)
for i := 0; i < dim*rows/8; i++ {
builder.ValueBuilder().(*array.Uint8Builder).Append(uint8(i))
rowBytes := dim * 2
for i := 0; i < rowBytes*rows; i++ {
builder.ValueBuilder().(*array.Uint8Builder).Append(uint8(i % 256))
}
for i := 0; i < rows; i++ {
offsets = append(offsets, int32(dim*i/8))
offsets = append(offsets, int32(rowBytes*i))
valid = append(valid, true)
}
builder.AppendValues(offsets, valid)
@ -348,7 +280,7 @@ func buildArrayData(dataType, elementType schemapb.DataType, dim, rows int, isBi
offsets = append(offsets, int32(index))
valid = append(valid, true)
}
switch elementType {
switch elemType {
case schemapb.DataType_Bool:
builder := array.NewListBuilder(mem, &arrow.BooleanType{})
valueBuilder := builder.ValueBuilder().(*array.BooleanBuilder)
@ -424,7 +356,10 @@ func GenerateParquetFile(filePath string, schema *schemapb.CollectionSchema, num
return err
}
pqSchema := convertMilvusSchemaToArrowSchema(schema)
pqSchema, err := pq.ConvertToArrowSchema(schema)
if err != nil {
return err
}
fw, err := pqarrow.NewFileWriter(pqSchema, w, parquet.NewWriterProperties(parquet.WithMaxRowGroupLength(int64(numRows))), pqarrow.DefaultWriterProps())
if err != nil {
return err
@ -436,7 +371,7 @@ func GenerateParquetFile(filePath string, schema *schemapb.CollectionSchema, num
if field.GetIsPrimaryKey() && field.GetAutoID() {
continue
}
columnData := buildArrayData(field.DataType, field.ElementType, dim, numRows, field.Name == "FieldBinaryVector2")
columnData := buildArrayData(field.DataType, field.ElementType, dim, numRows)
columns = append(columns, columnData)
}
recordBatch := array.NewRecord(pqSchema, columns, int64(numRows))
@ -542,10 +477,14 @@ func GenerateNumpyFile(filePath string, rowCount int, dType schemapb.DataType) e
return err
}
case schemapb.DataType_BinaryVector:
binVecData := make([]byte, 0)
total := rowCount * dim / 8
for i := 0; i < total; i++ {
binVecData = append(binVecData, byte(i%256))
const rowBytes = dim / 8
binVecData := make([][rowBytes]byte, 0, rowCount)
for i := 0; i < rowCount; i++ {
vec := [rowBytes]byte{}
for j := 0; j < rowBytes; j++ {
vec[j] = byte((i + j) % 256)
}
binVecData = append(binVecData, vec)
}
err := writeFn(filePath, binVecData)
if err != nil {
@ -556,7 +495,7 @@ func GenerateNumpyFile(filePath string, rowCount int, dType schemapb.DataType) e
for i := 0; i < rowCount; i++ {
vec := [dim]float32{}
for j := 0; j < dim; j++ {
vec[j] = 1.1
vec[j] = rand.Float32()
}
data = append(data, vec)
}
@ -564,14 +503,17 @@ func GenerateNumpyFile(filePath string, rowCount int, dType schemapb.DataType) e
if err != nil {
return err
}
case schemapb.DataType_Float16Vector:
total := int64(rowCount) * dim * 2
float16VecData := make([]byte, total)
_, err := rand2.Read(float16VecData)
if err != nil {
return err
case schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
const rowBytes = dim * 2
data := make([][rowBytes]byte, 0, rowCount)
for i := 0; i < rowCount; i++ {
vec := [rowBytes]byte{}
for j := 0; j < rowBytes; j++ {
vec[j] = byte(rand.Uint32() % 256)
}
data = append(data, vec)
}
err = writeFn(filePath, float16VecData)
err := writeFn(filePath, data)
if err != nil {
return err
}
@ -628,18 +570,19 @@ func GenerateJSONFile(t *testing.T, filePath string, schema *schemapb.Collection
if fieldIDToField[fieldID].GetAutoID() {
continue
}
if dataType == schemapb.DataType_Array {
switch dataType {
case schemapb.DataType_Array:
data[fieldID] = v.GetRow(i).(*schemapb.ScalarField).GetIntData().GetData()
} else if dataType == schemapb.DataType_JSON {
case schemapb.DataType_JSON:
data[fieldID] = string(v.GetRow(i).([]byte))
} else if dataType == schemapb.DataType_BinaryVector || dataType == schemapb.DataType_Float16Vector {
case schemapb.DataType_BinaryVector, schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
bytes := v.GetRow(i).([]byte)
ints := make([]int, 0, len(bytes))
for _, b := range bytes {
ints = append(ints, int(b))
}
data[fieldID] = ints
} else {
default:
data[fieldID] = v.GetRow(i)
}
}

23
tests/integration/util_query.go
View File

@ -256,24 +256,23 @@ func constructPlaceholderGroup(nq, dim int, vectorType schemapb.DataType) *commo
}
values = append(values, ret)
}
// case schemapb.DataType_BFloat16Vector:
// placeholderType = commonpb.PlaceholderType_BFloat16Vector
// for i := 0; i < nq; i++ {
// total := dim * 2
// ret := make([]byte, total)
// _, err := rand.Read(ret)
// if err != nil {
// panic(err)
// }
// values = append(values, ret)
// }
case schemapb.DataType_BFloat16Vector:
placeholderType = commonpb.PlaceholderType_BFloat16Vector
for i := 0; i < nq; i++ {
total := dim * 2
ret := make([]byte, total)
_, err := rand.Read(ret)
if err != nil {
panic(err)
}
values = append(values, ret)
}
case schemapb.DataType_SparseFloatVector:
// for sparse, all query rows are encoded in a single byte array
values = make([][]byte, 0, 1)
placeholderType = commonpb.PlaceholderType_SparseFloatVector
sparseVecs := GenerateSparseFloatArray(nq)
values = append(values, sparseVecs.Contents...)
default:
panic("invalid vector data type")
}