read by count

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>

internal/util/importutilv2/numpy/reader.go

@@ -26,65 +26,57 @@
 )
 
 type Reader struct {
-	schema     *schemapb.CollectionSchema
-	bufferSize int
-	frs        map[int64]*FieldReader // fieldID -> FieldReader
+	schema *schemapb.CollectionSchema
+	count  int64
+	frs    map[int64]*FieldReader // fieldID -> FieldReader
 }
 
 func NewReader(schema *schemapb.CollectionSchema, readers map[int64]io.Reader, bufferSize int) (*Reader, error) {
 	fields := lo.KeyBy(schema.GetFields(), func(field *schemapb.FieldSchema) int64 {
 		return field.GetFieldID()
 	})
+	count, err := calcRowCount(bufferSize, schema)
+	if err != nil {
+		return nil, err
+	}
 	crs := make(map[int64]*FieldReader)
 	for fieldID, r := range readers {
 		cr, err := NewFieldReader(r, fields[fieldID])
 		if err != nil {
 			return nil, err
 		}
 		crs[fieldID] = cr
 	}
 	return &Reader{
-		schema:     schema,
-		bufferSize: bufferSize,
-		frs:        crs,
+		schema: schema,
+		count:  count,
+		frs:    crs,
 	}, nil
 }
 
 func (r *Reader) Read() (*storage.InsertData, error) {
 	insertData, err := storage.NewInsertData(r.schema)
 	if err != nil {
 		return nil, err
 	}
-OUTER:
-	for {
-		for fieldID, cr := range r.frs {
-			data, err := cr.Next(1)
-			if err != nil {
-				return nil, err
-			}
-			if data == nil {
-				break OUTER
-			}
-			err = insertData.Data[fieldID].AppendRows(data)
-			if err != nil {
-				return nil, err
-			}
-		}
-		if insertData.GetMemorySize() >= r.bufferSize {
-			break
+	for fieldID, cr := range r.frs {
+		data, err := cr.Next(r.count)
+		if err != nil {
+			return nil, err
 		}
-	}
-
-	for fieldID := range r.frs {
-		if insertData.Data[fieldID].RowNum() == 0 {
+		if data == nil {
 			return nil, nil
 		}
+		err = insertData.Data[fieldID].AppendRows(data)
+		if err != nil {
+			return nil, err
+		}
 	}
 	return insertData, nil
 }
 
 func (r *Reader) Close() {
 	for _, cr := range r.frs {
 		cr.Close()
 	}
 }

internal/util/importutilv2/numpy/reader_test.go

@@ -182,6 +182,12 @@ func (suite *ReaderSuite) run(dt schemapb.DataType) {
 				Name:         "pk",
 				IsPrimaryKey: true,
 				DataType:     suite.pkDataType,
+				TypeParams: []*commonpb.KeyValuePair{
+					{
+						Key:   "max_length",
+						Value: "256",
+					},
+				},
 			},
 			{
 				FieldID:  101,
@@ -199,6 +205,12 @@ func (suite *ReaderSuite) run(dt schemapb.DataType) {
 				Name:        dt.String(),
 				DataType:    dt,
 				ElementType: schemapb.DataType_Int32,
+				TypeParams: []*commonpb.KeyValuePair{
+					{
+						Key:   "max_length",
+						Value: "256",
+					},
+				},
 			},
 		},
 	}

internal/util/importutilv2/numpy/util.go

@@ -30,6 +30,7 @@
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
 	"github.com/milvus-io/milvus/pkg/util/merr"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
 
 var (
@@ -42,48 +43,48 @@
 func stringLen(dtype string) (int, bool, error) {
 	var utf bool
 	switch {
 	case reStrPre.MatchString(dtype), reStrPost.MatchString(dtype):
 		utf = false
 	case reUniPre.MatchString(dtype), reUniPost.MatchString(dtype):
 		utf = true
 	}
 
 	if m := reStrPre.FindStringSubmatch(dtype); m != nil {
 		v, err := strconv.Atoi(m[1])
 		if err != nil {
 			return 0, false, err
 		}
 		return v, utf, nil
 	}
 	if m := reStrPost.FindStringSubmatch(dtype); m != nil {
 		v, err := strconv.Atoi(m[1])
 		if err != nil {
 			return 0, false, err
 		}
 		return v, utf, nil
 	}
 	if m := reUniPre.FindStringSubmatch(dtype); m != nil {
 		v, err := strconv.Atoi(m[1])
 		if err != nil {
 			return 0, false, err
 		}
 		return v, utf, nil
 	}
 	if m := reUniPost.FindStringSubmatch(dtype); m != nil {
 		v, err := strconv.Atoi(m[1])
 		if err != nil {
 			return 0, false, err
 		}
 		return v, utf, nil
 	}
 
 	return 0, false, merr.WrapErrImportFailed(fmt.Sprintf("dtype '%s' of numpy file is not varchar data type", dtype))
 }
 
 func decodeUtf32(src []byte, order binary.ByteOrder) (string, error) {
 	if len(src)%4 != 0 {
 		return "", merr.WrapErrImportFailed(fmt.Sprintf("invalid utf32 bytes length %d, the byte array length should be multiple of 4", len(src)))
 	}
 
 	var str string
 	for len(src) > 0 {
@@ -96,12 +97,12 @@
 				isUtf16 = true
 			}
 			lowbytesPosition = 0
 		} else {
 			if src[0] == 0 && src[1] == 0 {
 				isUtf16 = true
 			}
 			lowbytesPosition = 2
 			uOrder = unicode.BigEndian
 		}
 
 		if isUtf16 {
@@ -111,26 +112,26 @@
 				decoder := unicode.UTF16(uOrder, unicode.IgnoreBOM).NewDecoder()
 				res, err := decoder.Bytes(src[lowbytesPosition : lowbytesPosition+2])
 				if err != nil {
 					return "", merr.WrapErrImportFailed(fmt.Sprintf("failed to decode utf32 binary bytes, error: %v", err))
 				}
 				str += string(res)
 			}
 		} else {
 			// convert the 4 bytes to a unicode and encode to utf8
 			// Golang strongly opposes utf32 coding, this kind of encoding has been excluded from standard lib
 			var x uint32
 			if order == binary.LittleEndian {
 				x = uint32(src[3])<<24 | uint32(src[2])<<16 | uint32(src[1])<<8 | uint32(src[0])
 			} else {
 				x = uint32(src[0])<<24 | uint32(src[1])<<16 | uint32(src[2])<<8 | uint32(src[3])
 			}
 			r := rune(x)
 			utf8Code := make([]byte, 4)
 			utf8.EncodeRune(utf8Code, r)
 			if r == utf8.RuneError {
 				return "", merr.WrapErrImportFailed(fmt.Sprintf("failed to convert 4 bytes unicode %d to utf8 rune", x))
 			}
 			str += string(utf8Code)
 		}
 
 		src = src[4:]
@@ -164,69 +165,78 @@
 			// Note: JSON field and VARCHAR field are using string type numpy
 			return schemapb.DataType_VarChar, nil
 		}
 		return schemapb.DataType_None, merr.WrapErrImportFailed(
 			fmt.Sprintf("the numpy file dtype '%s' is not supported", typeStr))
 	}
 }
 
 func wrapElementTypeError(eleType schemapb.DataType, field *schemapb.FieldSchema) error {
 	return merr.WrapErrImportFailed(fmt.Sprintf("expected element type '%s' for field '%s', got type '%T'",
 		field.GetDataType().String(), field.GetName(), eleType))
 }
 
 func wrapDimError(actualDim int, expectDim int, field *schemapb.FieldSchema) error {
 	return merr.WrapErrImportFailed(fmt.Sprintf("expected dim '%d' for %s field '%s', got dim '%d'",
 		expectDim, field.GetDataType().String(), field.GetName(), actualDim))
 }
 
 func wrapShapeError(actualShape int, expectShape int, field *schemapb.FieldSchema) error {
 	return merr.WrapErrImportFailed(fmt.Sprintf("expected shape '%d' for %s field '%s', got shape '%d'",
 		expectShape, field.GetDataType().String(), field.GetName(), actualShape))
 }
 
 func validateHeader(npyReader *npy.Reader, field *schemapb.FieldSchema, dim int) error {
 	elementType, err := convertNumpyType(npyReader.Header.Descr.Type)
 	if err != nil {
 		return err
 	}
 	shape := npyReader.Header.Descr.Shape
 
 	switch field.GetDataType() {
 	case schemapb.DataType_FloatVector:
 		if elementType != schemapb.DataType_Float && elementType != schemapb.DataType_Double {
 			return wrapElementTypeError(elementType, field)
 		}
 		if len(shape) != 2 {
 			return wrapShapeError(len(shape), 2, field)
 		}
 		if shape[1] != dim {
 			return wrapDimError(shape[1], dim, field)
 		}
 	case schemapb.DataType_BinaryVector:
 		if elementType != schemapb.DataType_BinaryVector {
 			return wrapElementTypeError(elementType, field)
 		}
 		if len(shape) != 2 {
 			return wrapShapeError(len(shape), 2, field)
 		}
 		if shape[1] != dim/8 {
 			return wrapDimError(shape[1]*8, dim, field)
 		}
 	case schemapb.DataType_VarChar, schemapb.DataType_JSON:
 		if len(shape) != 1 {
 			return wrapShapeError(len(shape), 1, field)
 		}
 	case schemapb.DataType_None, schemapb.DataType_Array,
 		schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector:
 		return merr.WrapErrImportFailed(fmt.Sprintf("unsupported data type: %s", field.GetDataType().String()))
 
 	default:
 		if elementType != field.GetDataType() {
 			return wrapElementTypeError(elementType, field)
 		}
 		if len(shape) != 1 {
 			return wrapShapeError(len(shape), 1, field)
 		}
 	}
 	return nil
 }
+
+func calcRowCount(bufferSize int, schema *schemapb.CollectionSchema) (int64, error) {
+	sizePerRecord, err := typeutil.EstimateSizePerRecord(schema)
+	if err != nil {
+		return 0, err
+	}
+	rowCount := int64(bufferSize) / int64(sizePerRecord)
+	return rowCount, nil
+}