feat: storage v2 serde reader and writer (#39667)

related: https://github.com/milvus-io/milvus/issues/39173

---------

Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>

internal/storage/binlog_iterator_test.go

@@ -43,7 +43,7 @@ func generateTestSchema() *schemapb.CollectionSchema {
 		{FieldID: 16, Name: "varchar", DataType: schemapb.DataType_VarChar, IsPrimaryKey: true},
 		{FieldID: 17, Name: "string", DataType: schemapb.DataType_String},
 		{FieldID: 18, Name: "array", DataType: schemapb.DataType_Array},
-		{FieldID: 19, Name: "string", DataType: schemapb.DataType_JSON},
+		{FieldID: 19, Name: "json", DataType: schemapb.DataType_JSON},
 		{FieldID: 101, Name: "int32", DataType: schemapb.DataType_Int32},
 		{FieldID: 102, Name: "floatVector", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{
 			{Key: common.DimKey, Value: "8"},

internal/storage/schema.go

@@ -0,0 +1,38 @@
+package storage
+
+import (
+	"github.com/apache/arrow/go/v12/arrow"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
+	"github.com/milvus-io/milvus/pkg/util/merr"
+)
+
+func ConvertToArrowSchema(fields []*schemapb.FieldSchema) (*arrow.Schema, error) {
+	arrowFields := make([]arrow.Field, 0, len(fields))
+	for _, field := range fields {
+		if serdeMap[field.DataType].arrowType == nil {
+			return nil, merr.WrapErrParameterInvalidMsg("unknown field data type [%s] for field [%s]", field.DataType, field.GetName())
+		}
+		switch field.DataType {
+		case schemapb.DataType_BinaryVector, schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector,
+			schemapb.DataType_Int8Vector, schemapb.DataType_FloatVector:
+			dim, err := GetDimFromParams(field.TypeParams)
+			if err != nil {
+				return nil, merr.WrapErrParameterInvalidMsg("dim not found in field [%s] params", field.GetName())
+			}
+			arrowFields = append(arrowFields, arrow.Field{
+				Name:     field.GetName(),
+				Type:     serdeMap[field.DataType].arrowType(dim),
+				Nullable: field.GetNullable(),
+			})
+		default:
+			arrowFields = append(arrowFields, arrow.Field{
+				Name:     field.GetName(),
+				Type:     serdeMap[field.DataType].arrowType(0),
+				Nullable: field.GetNullable(),
+			})
+		}
+	}
+
+	return arrow.NewSchema(arrowFields, nil), nil
+}

internal/storage/schema_test.go

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package typeutil
+package storage
 
 import (
 	"testing"

internal/storage/serde_events.go

@@ -400,20 +400,20 @@ func ValueSerializer(v []*Value, fieldSchema []*schemapb.FieldSchema) (Record, e
 			}
 		}
 	}
-	arrays := make([]arrow.Array, len(types))
-	fields := make([]arrow.Field, len(types))
-	field2Col := make(map[FieldID]int, len(types))
-	i := 0
-	for fid, builder := range builders {
+	arrays := make([]arrow.Array, len(fieldSchema))
+	fields := make([]arrow.Field, len(fieldSchema))
+	field2Col := make(map[FieldID]int, len(fieldSchema))
+	for i, field := range fieldSchema {
+		builder := builders[field.FieldID]
 		arrays[i] = builder.NewArray()
 		builder.Release()
 		fields[i] = arrow.Field{
-			Name:     strconv.Itoa(int(fid)),
+			Name:     field.Name,
 			Type:     arrays[i].DataType(),
 			Nullable: true, // No nullable check here.
+			Metadata: arrow.NewMetadata([]string{"FieldID"}, []string{strconv.Itoa(int(field.FieldID))}),
 		}
-		field2Col[fid] = i
-		i++
+		field2Col[field.FieldID] = i
 	}
 	return newSimpleArrowRecord(array.NewRecord(arrow.NewSchema(fields, nil), arrays, int64(len(v))), field2Col), nil
 }

internal/storage/serde_events_v2.go

@@ -0,0 +1,214 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package storage
+
+import (
+	"fmt"
+	"io"
+
+	"github.com/apache/arrow/go/v12/arrow"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
+	"github.com/milvus-io/milvus/internal/storagev2/packed"
+	"github.com/milvus-io/milvus/pkg/common"
+	"github.com/milvus-io/milvus/pkg/util/merr"
+)
+
+type packedRecordReader struct {
+	reader *packed.PackedReader
+
+	bufferSize int
+	path       string
+	schema     *schemapb.CollectionSchema
+	r          *simpleArrowRecord
+	field2Col  map[FieldID]int
+}
+
+func (pr *packedRecordReader) Next() error {
+	if pr.reader == nil {
+		return io.EOF
+	}
+	rec, err := pr.reader.ReadNext()
+	if err != nil || rec == nil {
+		return io.EOF
+	}
+	pr.r = newSimpleArrowRecord(rec, pr.field2Col)
+	return nil
+}
+
+func (pr *packedRecordReader) Record() Record {
+	return pr.r
+}
+
+func (pr *packedRecordReader) Close() error {
+	if pr.reader != nil {
+		return pr.reader.Close()
+	}
+	return nil
+}
+
+func NewPackedRecordReader(path string, schema *schemapb.CollectionSchema, bufferSize int,
+) (*packedRecordReader, error) {
+	arrowSchema, err := ConvertToArrowSchema(schema.Fields)
+	if err != nil {
+		return nil, merr.WrapErrParameterInvalid("convert collection schema [%s] to arrow schema error: %s", schema.Name, err.Error())
+	}
+	reader, err := packed.NewPackedReader(path, arrowSchema, bufferSize)
+	if err != nil {
+		return nil, merr.WrapErrParameterInvalid("New binlog record packed reader error: %s", err.Error())
+	}
+	field2Col := make(map[FieldID]int)
+	for i, field := range schema.Fields {
+		field2Col[field.FieldID] = i
+	}
+	return &packedRecordReader{
+		reader:     reader,
+		schema:     schema,
+		bufferSize: bufferSize,
+		path:       path,
+		field2Col:  field2Col,
+	}, nil
+}
+
+func NewPackedDeserializeReader(path string, schema *schemapb.CollectionSchema,
+	bufferSize int, pkFieldID FieldID,
+) (*DeserializeReader[*Value], error) {
+	reader, err := NewPackedRecordReader(path, schema, bufferSize)
+	if err != nil {
+		return nil, err
+	}
+
+	return NewDeserializeReader(reader, func(r Record, v []*Value) error {
+		rec, ok := r.(*simpleArrowRecord)
+		if !ok {
+			return merr.WrapErrServiceInternal("can not cast to simple arrow record")
+		}
+
+		schema := reader.schema
+		numFields := len(schema.Fields)
+		for i := 0; i < rec.Len(); i++ {
+			if v[i] == nil {
+				v[i] = &Value{
+					Value: make(map[FieldID]interface{}, numFields),
+				}
+			}
+			value := v[i]
+			m := value.Value.(map[FieldID]interface{})
+			for _, field := range schema.Fields {
+				fieldID := field.FieldID
+				column := r.Column(fieldID)
+				if column.IsNull(i) {
+					m[fieldID] = nil
+				} else {
+					d, ok := serdeMap[field.DataType].deserialize(column, i)
+					if ok {
+						m[fieldID] = d
+					} else {
+						return merr.WrapErrServiceInternal(fmt.Sprintf("can not deserialize field [%s]", field.Name))
+					}
+				}
+			}
+
+			rowID, ok := m[common.RowIDField].(int64)
+			if !ok {
+				return merr.WrapErrIoKeyNotFound("no row id column found")
+			}
+			value.ID = rowID
+			value.Timestamp = m[common.TimeStampField].(int64)
+
+			pkCol := rec.field2Col[pkFieldID]
+			pk, err := GenPrimaryKeyByRawData(m[pkFieldID], schema.Fields[pkCol].DataType)
+			if err != nil {
+				return err
+			}
+
+			value.PK = pk
+			value.IsDeleted = false
+			value.Value = m
+		}
+		return nil
+	}), nil
+}
+
+var _ RecordWriter = (*packedRecordWriter)(nil)
+
+type packedRecordWriter struct {
+	writer *packed.PackedWriter
+
+	bufferSize int
+	path       string
+	schema     *arrow.Schema
+
+	numRows             int
+	writtenUncompressed uint64
+}
+
+func (pw *packedRecordWriter) Write(r Record) error {
+	rec, ok := r.(*simpleArrowRecord)
+	if !ok {
+		return merr.WrapErrServiceInternal("can not cast to simple arrow record")
+	}
+	pw.numRows += r.Len()
+	for _, arr := range rec.r.Columns() {
+		pw.writtenUncompressed += uint64(calculateArraySize(arr))
+	}
+	defer rec.Release()
+	return pw.writer.WriteRecordBatch(rec.r)
+}
+
+func (pw *packedRecordWriter) GetWrittenUncompressed() uint64 {
+	return pw.writtenUncompressed
+}
+
+func (pw *packedRecordWriter) Close() error {
+	if pw.writer != nil {
+		return pw.writer.Close()
+	}
+	return nil
+}
+
+func NewPackedRecordWriter(path string, schema *arrow.Schema, bufferSize int) (*packedRecordWriter, error) {
+	writer, err := packed.NewPackedWriter(path, schema, bufferSize)
+	if err != nil {
+		return nil, merr.WrapErrServiceInternal(
+			fmt.Sprintf("can not new packed record writer %s", err.Error()))
+	}
+	return &packedRecordWriter{
+		writer:     writer,
+		schema:     schema,
+		bufferSize: bufferSize,
+		path:       path,
+	}, nil
+}
+
+func NewPackedSerializeWriter(schema *schemapb.CollectionSchema, partitionID, segmentID UniqueID,
+	batchSize int, path string, bufferSize int,
+) (*SerializeWriter[*Value], error) {
+	arrowSchema, err := ConvertToArrowSchema(schema.Fields)
+	if err != nil {
+		return nil, merr.WrapErrServiceInternal(
+			fmt.Sprintf("can not convert collection schema %s to arrow schema: %s", schema.Name, err.Error()))
+	}
+	packedRecordWriter, err := NewPackedRecordWriter(path, arrowSchema, bufferSize)
+	if err != nil {
+		return nil, merr.WrapErrServiceInternal(
+			fmt.Sprintf("can not new packed record writer %s", err.Error()))
+	}
+	return NewSerializeRecordWriter[*Value](packedRecordWriter, func(v []*Value) (Record, error) {
+		return ValueSerializer(v, schema.Fields)
+	}, batchSize), nil
+}

internal/storage/serde_events_v2_test.go

@@ -0,0 +1,68 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package storage
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/milvus-io/milvus/pkg/common"
+)
+
+func TestPackedSerde(t *testing.T) {
+	t.Run("test binlog packed deserialize reader", func(t *testing.T) {
+		size := 10
+
+		blobs, err := generateTestData(size)
+		assert.NoError(t, err)
+
+		reader, err := NewBinlogDeserializeReader(blobs, common.RowIDField)
+		assert.NoError(t, err)
+		defer reader.Close()
+
+		path := "/tmp"
+		bufferSize := 10 * 1024 * 1024 // 10MB
+		schema := generateTestSchema()
+
+		writer, err := NewPackedSerializeWriter(schema, 0, 0, 7, path, bufferSize)
+		assert.NoError(t, err)
+
+		for i := 1; i <= size; i++ {
+			err = reader.Next()
+			assert.NoError(t, err)
+
+			value := reader.Value()
+			assertTestData(t, i, value)
+			err := writer.Write(value)
+			assert.NoError(t, err)
+		}
+		err = writer.Close()
+		assert.NoError(t, err)
+
+		reader, err = NewPackedDeserializeReader(path, schema, bufferSize, common.RowIDField)
+		assert.NoError(t, err)
+		defer reader.Close()
+
+		for i := 1; i <= size; i++ {
+			err = reader.Next()
+			assert.NoError(t, err)
+			value := reader.Value()
+			assertTestData(t, i, value)
+		}
+	})
+}

internal/storagev2/packed/packed_reader.go

@@ -37,7 +37,6 @@ func NewPackedReader(path string, schema *arrow.Schema, bufferSize int) (*Packed
 	var cas cdata.CArrowSchema
 	cdata.ExportArrowSchema(schema, &cas)
 	cSchema := (*C.struct_ArrowSchema)(unsafe.Pointer(&cas))
-
 	cPath := C.CString(path)
 	defer C.free(unsafe.Pointer(cPath))
 

internal/util/typeutil/schema.go

@@ -1,139 +0,0 @@
-package typeutil
-
-import (
-	"github.com/apache/arrow/go/v12/arrow"
-
-	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
-	"github.com/milvus-io/milvus/internal/storage"
-	"github.com/milvus-io/milvus/pkg/util/merr"
-)
-
-func ConvertToArrowSchema(fields []*schemapb.FieldSchema) (*arrow.Schema, error) {
-	arrowFields := make([]arrow.Field, 0, len(fields))
-	for _, field := range fields {
-		switch field.DataType {
-		case schemapb.DataType_Bool:
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: arrow.FixedWidthTypes.Boolean,
-			})
-		case schemapb.DataType_Int8:
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: arrow.PrimitiveTypes.Int8,
-			})
-		case schemapb.DataType_Int16:
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: arrow.PrimitiveTypes.Int16,
-			})
-		case schemapb.DataType_Int32:
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: arrow.PrimitiveTypes.Int32,
-			})
-		case schemapb.DataType_Int64:
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: arrow.PrimitiveTypes.Int64,
-			})
-		case schemapb.DataType_Float:
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: arrow.PrimitiveTypes.Float32,
-			})
-		case schemapb.DataType_Double:
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: arrow.PrimitiveTypes.Float64,
-			})
-		case schemapb.DataType_String, schemapb.DataType_VarChar:
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: arrow.BinaryTypes.String,
-			})
-		case schemapb.DataType_Array:
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: arrow.BinaryTypes.Binary,
-			})
-		case schemapb.DataType_JSON:
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: arrow.BinaryTypes.Binary,
-			})
-		case schemapb.DataType_BinaryVector:
-			dim, err := storage.GetDimFromParams(field.TypeParams)
-			if err != nil {
-				return nil, err
-			}
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: &arrow.FixedSizeBinaryType{ByteWidth: dim / 8},
-			})
-		case schemapb.DataType_FloatVector:
-			dim, err := storage.GetDimFromParams(field.TypeParams)
-			if err != nil {
-				return nil, err
-			}
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: &arrow.FixedSizeBinaryType{ByteWidth: dim * 4},
-			})
-		case schemapb.DataType_Float16Vector:
-			dim, err := storage.GetDimFromParams(field.TypeParams)
-			if err != nil {
-				return nil, err
-			}
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: &arrow.FixedSizeBinaryType{ByteWidth: dim * 2},
-			})
-		case schemapb.DataType_BFloat16Vector:
-			dim, err := storage.GetDimFromParams(field.TypeParams)
-			if err != nil {
-				return nil, err
-			}
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: &arrow.FixedSizeBinaryType{ByteWidth: dim * 2},
-			})
-		case schemapb.DataType_Int8Vector:
-			dim, err := storage.GetDimFromParams(field.TypeParams)
-			if err != nil {
-				return nil, err
-			}
-			arrowFields = append(arrowFields, arrow.Field{
-				Name: field.Name,
-				Type: &arrow.FixedSizeBinaryType{ByteWidth: dim},
-			})
-		default:
-			return nil, merr.WrapErrParameterInvalidMsg("unknown type %v", field.DataType.String())
-		}
-	}
-
-	return arrow.NewSchema(arrowFields, nil), nil
-}
-
-func convertToArrowType(dataType schemapb.DataType) (arrow.DataType, error) {
-	switch dataType {
-	case schemapb.DataType_Bool:
-		return arrow.FixedWidthTypes.Boolean, nil
-	case schemapb.DataType_Int8:
-		return arrow.PrimitiveTypes.Int8, nil
-	case schemapb.DataType_Int16:
-		return arrow.PrimitiveTypes.Int16, nil
-	case schemapb.DataType_Int32:
-		return arrow.PrimitiveTypes.Int32, nil
-	case schemapb.DataType_Int64:
-		return arrow.PrimitiveTypes.Int64, nil
-	case schemapb.DataType_Float:
-		return arrow.PrimitiveTypes.Float32, nil
-	case schemapb.DataType_Double:
-		return arrow.PrimitiveTypes.Float64, nil
-	case schemapb.DataType_String, schemapb.DataType_VarChar:
-		return arrow.BinaryTypes.String, nil
-	default:
-		return nil, merr.WrapErrParameterInvalidMsg("unknown type %v", dataType.String())
-	}
-}