enhance: Handle Float16Vector/BFloat16Vector numpy bulk insert as same as BinaryVector (#33760)

Issue: #22837

Signed-off-by: Cai Yudong <yudong.cai@zilliz.com>
pull/33766/head
Cai Yudong 2024-06-12 17:17:55 +08:00 committed by GitHub
parent 47b04ea167
commit 9d4535ce0b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 9 additions and 21 deletions

View File

@ -195,7 +195,7 @@ func (c *FieldReader) Next(count int64) (any, error) {
}
data = byteArr
c.readPosition += int(readCount)
case schemapb.DataType_BinaryVector:
case schemapb.DataType_BinaryVector, schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
data, err = ReadN[uint8](c.reader, c.order, readCount)
if err != nil {
return nil, err
@ -232,12 +232,6 @@ 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()))
}

View File

@ -199,6 +199,7 @@ func (s *BulkInsertSuite) run() {
searchResult, err := c.Proxy.Search(ctx, searchReq)
s.NoError(err)
s.Equal(commonpb.ErrorCode_Success, searchResult.GetStatus().GetErrorCode())
// s.Equal(nq*topk, len(searchResult.GetResults().GetScores()))
}
func (s *BulkInsertSuite) TestMultiFileTypes() {

View File

@ -31,6 +31,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/pkg/common"
"github.com/milvus-io/milvus/pkg/util/testutils"
)
const (
@ -247,25 +248,17 @@ func constructPlaceholderGroup(nq, dim int, vectorType schemapb.DataType) *commo
}
case schemapb.DataType_Float16Vector:
placeholderType = commonpb.PlaceholderType_Float16Vector
data := testutils.GenerateFloat16Vectors(nq, dim)
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)
rowBytes := dim * 2
values = append(values, data[rowBytes*i:rowBytes*(i+1)])
}
case schemapb.DataType_BFloat16Vector:
placeholderType = commonpb.PlaceholderType_BFloat16Vector
data := testutils.GenerateBFloat16Vectors(nq, dim)
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)
rowBytes := dim * 2
values = append(values, data[rowBytes*i:rowBytes*(i+1)])
}
case schemapb.DataType_SparseFloatVector:
// for sparse, all query rows are encoded in a single byte array