enhance: Handle Float16Vector/BFloat16Vector numpy bulk insert as sam…

…e as BinaryVector (milvus-io#33760) Issue: milvus-io#22837 Signed-off-by: Cai Yudong <yudong.cai@zilliz.com>
cydrain · Jun 12, 2024 · e5669f1 · e5669f1
1 parent eeba851
commit e5669f1
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 21 deletions.
diff --git a/internal/util/importutilv2/numpy/field_reader.go b/internal/util/importutilv2/numpy/field_reader.go
@@ -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()))
 	}

diff --git a/tests/integration/import/import_test.go b/tests/integration/import/import_test.go
@@ -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() {

diff --git a/tests/integration/util_query.go b/tests/integration/util_query.go
@@ -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