mirror of https://github.com/milvus-io/milvus.git
437 lines
13 KiB
Go
437 lines
13 KiB
Go
package importutil
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"errors"
|
|
"os"
|
|
"path"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/milvus-io/milvus/internal/allocator"
|
|
"github.com/milvus-io/milvus/internal/common"
|
|
"github.com/milvus-io/milvus/internal/log"
|
|
"github.com/milvus-io/milvus/internal/proto/schemapb"
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
"github.com/milvus-io/milvus/internal/util/typeutil"
|
|
"go.uber.org/zap"
|
|
"go.uber.org/zap/zapcore"
|
|
)
|
|
|
|
const (
|
|
JSONFileExt = ".json"
|
|
NumpyFileExt = ".npy"
|
|
)
|
|
|
|
type ImportWrapper struct {
|
|
ctx context.Context // for canceling parse process
|
|
cancel context.CancelFunc // for canceling parse process
|
|
collectionSchema *schemapb.CollectionSchema // collection schema
|
|
shardNum int32 // sharding number of the collection
|
|
segmentSize int32 // maximum size of a segment in MB
|
|
rowIDAllocator *allocator.IDAllocator // autoid allocator
|
|
|
|
callFlushFunc func(fields map[string]storage.FieldData) error // call back function to flush a segment
|
|
}
|
|
|
|
func NewImportWrapper(ctx context.Context, collectionSchema *schemapb.CollectionSchema, shardNum int32, segmentSize int32,
|
|
idAlloc *allocator.IDAllocator, flushFunc func(fields map[string]storage.FieldData) error) *ImportWrapper {
|
|
if collectionSchema == nil {
|
|
log.Error("import error: collection schema is nil")
|
|
return nil
|
|
}
|
|
|
|
// ignore the RowID field and Timestamp field
|
|
realSchema := &schemapb.CollectionSchema{
|
|
Name: collectionSchema.GetName(),
|
|
Description: collectionSchema.GetDescription(),
|
|
AutoID: collectionSchema.GetAutoID(),
|
|
Fields: make([]*schemapb.FieldSchema, 0),
|
|
}
|
|
for i := 0; i < len(collectionSchema.Fields); i++ {
|
|
schema := collectionSchema.Fields[i]
|
|
if schema.GetName() == common.RowIDFieldName || schema.GetName() == common.TimeStampFieldName {
|
|
continue
|
|
}
|
|
realSchema.Fields = append(realSchema.Fields, schema)
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(ctx)
|
|
|
|
wrapper := &ImportWrapper{
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
collectionSchema: realSchema,
|
|
shardNum: shardNum,
|
|
segmentSize: segmentSize,
|
|
rowIDAllocator: idAlloc,
|
|
callFlushFunc: flushFunc,
|
|
}
|
|
|
|
return wrapper
|
|
}
|
|
|
|
// this method can be used to cancel parse process
|
|
func (p *ImportWrapper) Cancel() error {
|
|
p.cancel()
|
|
return nil
|
|
}
|
|
|
|
func (p *ImportWrapper) printFieldsDataInfo(fieldsData map[string]storage.FieldData, msg string, files []string) {
|
|
stats := make([]zapcore.Field, 0)
|
|
for k, v := range fieldsData {
|
|
stats = append(stats, zap.Int(k, v.RowNum()))
|
|
}
|
|
|
|
if len(files) > 0 {
|
|
stats = append(stats, zap.Any("files", files))
|
|
}
|
|
log.Debug(msg, stats...)
|
|
}
|
|
|
|
func getFileNameAndExt(filePath string) (string, string) {
|
|
fileName := path.Base(filePath)
|
|
fileType := path.Ext(fileName)
|
|
fileNameWithoutExt := strings.TrimSuffix(fileName, fileType)
|
|
return fileNameWithoutExt, fileType
|
|
}
|
|
|
|
// import process entry
|
|
// filePath and rowBased are from ImportTask
|
|
// if onlyValidate is true, this process only do validation, no data generated, callFlushFunc will not be called
|
|
func (p *ImportWrapper) Import(filePaths []string, rowBased bool, onlyValidate bool) error {
|
|
if rowBased {
|
|
// parse and consume row-based files
|
|
// for row-based files, the JSONRowConsumer will generate autoid for primary key, and split rows into segments
|
|
// according to shard number, so the callFlushFunc will be called in the JSONRowConsumer
|
|
for i := 0; i < len(filePaths); i++ {
|
|
filePath := filePaths[i]
|
|
_, fileType := getFileNameAndExt(filePath)
|
|
log.Debug("imprort wrapper: row-based file ", zap.Any("filePath", filePath), zap.Any("fileType", fileType))
|
|
|
|
if fileType == JSONFileExt {
|
|
err := func() error {
|
|
file, err := os.Open(filePath)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer file.Close()
|
|
|
|
reader := bufio.NewReader(file)
|
|
parser := NewJSONParser(p.ctx, p.collectionSchema)
|
|
var consumer *JSONRowConsumer
|
|
if !onlyValidate {
|
|
flushFunc := func(fields map[string]storage.FieldData) error {
|
|
p.printFieldsDataInfo(fields, "import wrapper: prepare to flush segment", filePaths)
|
|
return p.callFlushFunc(fields)
|
|
}
|
|
consumer = NewJSONRowConsumer(p.collectionSchema, p.rowIDAllocator, p.shardNum, p.segmentSize, flushFunc)
|
|
}
|
|
validator := NewJSONRowValidator(p.collectionSchema, consumer)
|
|
|
|
err = parser.ParseRows(reader, validator)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}()
|
|
|
|
if err != nil {
|
|
log.Error("imprort error: "+err.Error(), zap.String("filePath", filePath))
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
// parse and consume row-based files
|
|
// for column-based files, the XXXColumnConsumer only output map[string]storage.FieldData
|
|
// after all columns are parsed/consumed, we need to combine map[string]storage.FieldData into one
|
|
// and use splitFieldsData() to split fields data into segments according to shard number
|
|
fieldsData := initSegmentData(p.collectionSchema)
|
|
rowCount := 0
|
|
|
|
// function to combine column data into fieldsData
|
|
combineFunc := func(fields map[string]storage.FieldData) error {
|
|
if len(fields) == 0 {
|
|
return nil
|
|
}
|
|
|
|
p.printFieldsDataInfo(fields, "imprort wrapper: combine field data", nil)
|
|
|
|
fieldNames := make([]string, 0)
|
|
for k, v := range fields {
|
|
// ignore 0 row field
|
|
if v.RowNum() == 0 {
|
|
continue
|
|
}
|
|
|
|
// each column should be only combined once
|
|
data, ok := fieldsData[k]
|
|
if ok && data.RowNum() > 0 {
|
|
return errors.New("the field " + k + " is duplicated")
|
|
}
|
|
|
|
// check the row count. only count non-zero row fields
|
|
if rowCount > 0 && rowCount != v.RowNum() {
|
|
return errors.New("the field " + k + " row count " + strconv.Itoa(v.RowNum()) + " doesn't equal " + strconv.Itoa(rowCount))
|
|
}
|
|
rowCount = v.RowNum()
|
|
|
|
// assign column data to fieldsData
|
|
fieldsData[k] = v
|
|
fieldNames = append(fieldNames, k)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// parse/validate/consume data
|
|
for i := 0; i < len(filePaths); i++ {
|
|
filePath := filePaths[i]
|
|
fileName, fileType := getFileNameAndExt(filePath)
|
|
log.Debug("imprort wrapper: column-based file ", zap.Any("filePath", filePath), zap.Any("fileType", fileType))
|
|
|
|
if fileType == JSONFileExt {
|
|
err := func() error {
|
|
file, err := os.Open(filePath)
|
|
if err != nil {
|
|
log.Error("imprort error: "+err.Error(), zap.String("filePath", filePath))
|
|
return err
|
|
}
|
|
defer file.Close()
|
|
|
|
reader := bufio.NewReader(file)
|
|
parser := NewJSONParser(p.ctx, p.collectionSchema)
|
|
var consumer *JSONColumnConsumer
|
|
if !onlyValidate {
|
|
consumer = NewJSONColumnConsumer(p.collectionSchema, combineFunc)
|
|
}
|
|
validator := NewJSONColumnValidator(p.collectionSchema, consumer)
|
|
|
|
err = parser.ParseColumns(reader, validator)
|
|
if err != nil {
|
|
log.Error("imprort error: "+err.Error(), zap.String("filePath", filePath))
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}()
|
|
|
|
if err != nil {
|
|
log.Error("imprort error: "+err.Error(), zap.String("filePath", filePath))
|
|
return err
|
|
}
|
|
} else if fileType == NumpyFileExt {
|
|
file, err := os.Open(filePath)
|
|
if err != nil {
|
|
log.Error("imprort error: "+err.Error(), zap.String("filePath", filePath))
|
|
return err
|
|
}
|
|
defer file.Close()
|
|
|
|
// the numpy parser return a storage.FieldData, here construct a map[string]storage.FieldData to combine
|
|
flushFunc := func(field storage.FieldData) error {
|
|
fields := make(map[string]storage.FieldData)
|
|
fields[fileName] = field
|
|
combineFunc(fields)
|
|
return nil
|
|
}
|
|
|
|
// for numpy file, we say the file name(without extension) is the filed name
|
|
parser := NewNumpyParser(p.ctx, p.collectionSchema, flushFunc)
|
|
err = parser.Parse(file, fileName, onlyValidate)
|
|
if err != nil {
|
|
log.Error("imprort error: "+err.Error(), zap.String("filePath", filePath))
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
// split fields data into segments
|
|
err := p.splitFieldsData(fieldsData, filePaths)
|
|
if err != nil {
|
|
log.Error("imprort error: " + err.Error())
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (p *ImportWrapper) appendFunc(schema *schemapb.FieldSchema) func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
switch schema.DataType {
|
|
case schemapb.DataType_Bool:
|
|
return func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
arr := target.(*storage.BoolFieldData)
|
|
arr.Data = append(arr.Data, src.GetRow(n).(bool))
|
|
arr.NumRows[0]++
|
|
return nil
|
|
}
|
|
case schemapb.DataType_Float:
|
|
return func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
arr := target.(*storage.FloatFieldData)
|
|
arr.Data = append(arr.Data, src.GetRow(n).(float32))
|
|
arr.NumRows[0]++
|
|
return nil
|
|
}
|
|
case schemapb.DataType_Double:
|
|
return func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
arr := target.(*storage.DoubleFieldData)
|
|
arr.Data = append(arr.Data, src.GetRow(n).(float64))
|
|
arr.NumRows[0]++
|
|
return nil
|
|
}
|
|
case schemapb.DataType_Int8:
|
|
return func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
arr := target.(*storage.Int8FieldData)
|
|
arr.Data = append(arr.Data, src.GetRow(n).(int8))
|
|
arr.NumRows[0]++
|
|
return nil
|
|
}
|
|
case schemapb.DataType_Int16:
|
|
return func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
arr := target.(*storage.Int16FieldData)
|
|
arr.Data = append(arr.Data, src.GetRow(n).(int16))
|
|
arr.NumRows[0]++
|
|
return nil
|
|
}
|
|
case schemapb.DataType_Int32:
|
|
return func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
arr := target.(*storage.Int32FieldData)
|
|
arr.Data = append(arr.Data, src.GetRow(n).(int32))
|
|
arr.NumRows[0]++
|
|
return nil
|
|
}
|
|
case schemapb.DataType_Int64:
|
|
return func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
arr := target.(*storage.Int64FieldData)
|
|
arr.Data = append(arr.Data, src.GetRow(n).(int64))
|
|
arr.NumRows[0]++
|
|
return nil
|
|
}
|
|
case schemapb.DataType_BinaryVector:
|
|
return func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
arr := target.(*storage.BinaryVectorFieldData)
|
|
arr.Data = append(arr.Data, src.GetRow(n).([]byte)...)
|
|
arr.NumRows[0]++
|
|
return nil
|
|
}
|
|
case schemapb.DataType_FloatVector:
|
|
return func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
arr := target.(*storage.FloatVectorFieldData)
|
|
arr.Data = append(arr.Data, src.GetRow(n).([]float32)...)
|
|
arr.NumRows[0]++
|
|
return nil
|
|
}
|
|
case schemapb.DataType_String:
|
|
return func(src storage.FieldData, n int, target storage.FieldData) error {
|
|
arr := target.(*storage.StringFieldData)
|
|
arr.Data = append(arr.Data, src.GetRow(n).(string))
|
|
return nil
|
|
}
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (p *ImportWrapper) splitFieldsData(fieldsData map[string]storage.FieldData, files []string) error {
|
|
if len(fieldsData) == 0 {
|
|
return errors.New("imprort error: fields data is empty")
|
|
}
|
|
|
|
var primaryKey *schemapb.FieldSchema
|
|
for i := 0; i < len(p.collectionSchema.Fields); i++ {
|
|
schema := p.collectionSchema.Fields[i]
|
|
if schema.GetIsPrimaryKey() {
|
|
primaryKey = schema
|
|
} else {
|
|
_, ok := fieldsData[schema.GetName()]
|
|
if !ok {
|
|
return errors.New("imprort error: field " + schema.GetName() + " not provided")
|
|
}
|
|
}
|
|
}
|
|
if primaryKey == nil {
|
|
return errors.New("imprort error: primary key field is not found")
|
|
}
|
|
|
|
rowCount := 0
|
|
for _, v := range fieldsData {
|
|
rowCount = v.RowNum()
|
|
break
|
|
}
|
|
|
|
primaryData, ok := fieldsData[primaryKey.GetName()]
|
|
if !ok {
|
|
// generate auto id for primary key
|
|
if primaryKey.GetAutoID() {
|
|
var rowIDBegin typeutil.UniqueID
|
|
var rowIDEnd typeutil.UniqueID
|
|
rowIDBegin, rowIDEnd, _ = p.rowIDAllocator.Alloc(uint32(rowCount))
|
|
|
|
primaryDataArr := primaryData.(*storage.Int64FieldData)
|
|
for i := rowIDBegin; i < rowIDEnd; i++ {
|
|
primaryDataArr.Data = append(primaryDataArr.Data, rowIDBegin+i)
|
|
}
|
|
}
|
|
}
|
|
|
|
if primaryData.RowNum() <= 0 {
|
|
return errors.New("imprort error: primary key " + primaryKey.GetName() + " not provided")
|
|
}
|
|
|
|
// prepare segemnts
|
|
segmentsData := make([]map[string]storage.FieldData, 0, p.shardNum)
|
|
for i := 0; i < int(p.shardNum); i++ {
|
|
segmentData := initSegmentData(p.collectionSchema)
|
|
if segmentData == nil {
|
|
return nil
|
|
}
|
|
segmentsData = append(segmentsData, segmentData)
|
|
}
|
|
|
|
// prepare append functions
|
|
appendFunctions := make(map[string]func(src storage.FieldData, n int, target storage.FieldData) error)
|
|
for i := 0; i < len(p.collectionSchema.Fields); i++ {
|
|
schema := p.collectionSchema.Fields[i]
|
|
appendFunc := p.appendFunc(schema)
|
|
if appendFunc == nil {
|
|
return errors.New("imprort error: unsupported field data type")
|
|
}
|
|
appendFunctions[schema.GetName()] = appendFunc
|
|
}
|
|
|
|
// split data into segments
|
|
for i := 0; i < rowCount; i++ {
|
|
id := primaryData.GetRow(i).(int64)
|
|
// hash to a shard number
|
|
hash, _ := typeutil.Hash32Int64(id)
|
|
shard := hash % uint32(p.shardNum)
|
|
|
|
for k := 0; k < len(p.collectionSchema.Fields); k++ {
|
|
schema := p.collectionSchema.Fields[k]
|
|
srcData := fieldsData[schema.GetName()]
|
|
targetData := segmentsData[shard][schema.GetName()]
|
|
appendFunc := appendFunctions[schema.GetName()]
|
|
err := appendFunc(srcData, i, targetData)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
// call flush function
|
|
for i := 0; i < int(p.shardNum); i++ {
|
|
segmentData := segmentsData[i]
|
|
p.printFieldsDataInfo(segmentData, "import wrapper: prepare to flush segment", files)
|
|
err := p.callFlushFunc(segmentData)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|