mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-28 22:45:26 +08:00
5e525eb3bf
This PR improves the robustness of object storage operations by retrying both explicit throttling errors (e.g. HTTP 429, SlowDown, ServerBusy). These errors commonly occur under high concurrency and are typically recoverable with bounded retries. issue: https://github.com/milvus-io/milvus/issues/44772 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Configurable retry support for reads from object storage and improved mapping of transient/rate-limit errors. * Added a retryable reader wrapper used by CSV/JSON/Parquet/Numpy import paths. * **Configuration** * New parameter to control storage read retry attempts. * **Tests** * Expanded unit tests covering error mapping and retry behaviors across storage backends. * Standardized mock readers and test initialization to simplify test setups. <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub> <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
513 lines
15 KiB
Go
513 lines
15 KiB
Go
// 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 json
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"math"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
"github.com/milvus-io/milvus/internal/mocks"
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
importcommon "github.com/milvus-io/milvus/internal/util/importutilv2/common"
|
|
"github.com/milvus-io/milvus/internal/util/nullutil"
|
|
"github.com/milvus-io/milvus/internal/util/testutil"
|
|
"github.com/milvus-io/milvus/pkg/v2/common"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
|
|
)
|
|
|
|
func init() {
|
|
paramtable.Init()
|
|
}
|
|
|
|
type ReaderSuite struct {
|
|
suite.Suite
|
|
|
|
numRows int
|
|
pkDataType schemapb.DataType
|
|
vecDataType schemapb.DataType
|
|
}
|
|
|
|
func (suite *ReaderSuite) SetupTest() {
|
|
// default suite params
|
|
suite.numRows = 100
|
|
suite.pkDataType = schemapb.DataType_Int64
|
|
suite.vecDataType = schemapb.DataType_FloatVector
|
|
}
|
|
|
|
func (suite *ReaderSuite) run(dataType schemapb.DataType, elemType schemapb.DataType, nullable bool) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "pk",
|
|
IsPrimaryKey: true,
|
|
DataType: suite.pkDataType,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.MaxLengthKey,
|
|
Value: "128",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "vec",
|
|
DataType: suite.vecDataType,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.DimKey,
|
|
Value: "8",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 102,
|
|
Name: dataType.String(),
|
|
DataType: dataType,
|
|
ElementType: elemType,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.MaxLengthKey,
|
|
Value: "128",
|
|
},
|
|
{
|
|
Key: common.MaxCapacityKey,
|
|
Value: "128",
|
|
},
|
|
},
|
|
Nullable: nullable,
|
|
},
|
|
},
|
|
}
|
|
|
|
if dataType == schemapb.DataType_VarChar {
|
|
// Add a BM25 function if data type is VarChar
|
|
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
|
|
FieldID: 103,
|
|
Name: "sparse",
|
|
DataType: schemapb.DataType_SparseFloatVector,
|
|
IsFunctionOutput: true,
|
|
})
|
|
schema.Functions = append(schema.Functions, &schemapb.FunctionSchema{
|
|
Id: 1000,
|
|
Name: "bm25",
|
|
Type: schemapb.FunctionType_BM25,
|
|
InputFieldIds: []int64{102},
|
|
InputFieldNames: []string{dataType.String()},
|
|
OutputFieldIds: []int64{103},
|
|
OutputFieldNames: []string{"sparse"},
|
|
})
|
|
}
|
|
|
|
insertData, err := testutil.CreateInsertData(schema, suite.numRows)
|
|
suite.NoError(err)
|
|
|
|
rows, err := testutil.CreateInsertDataRowsForJSON(schema, insertData)
|
|
suite.NoError(err)
|
|
|
|
jsonBytes, err := json.Marshal(rows)
|
|
suite.NoError(err)
|
|
|
|
cm := mocks.NewChunkManager(suite.T())
|
|
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
|
|
r := importcommon.NewMockReader(string(jsonBytes))
|
|
return r, nil
|
|
})
|
|
cm.EXPECT().Size(mock.Anything, "mockPath").Return(128, nil)
|
|
reader, err := NewReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
|
|
suite.NoError(err)
|
|
suite.NotNil(reader)
|
|
defer reader.Close()
|
|
|
|
size, err := reader.Size()
|
|
suite.NoError(err)
|
|
suite.Equal(int64(128), size)
|
|
size2, err := reader.Size() // size is cached
|
|
suite.NoError(err)
|
|
suite.Equal(size, size2)
|
|
|
|
checkFn := func(actualInsertData *storage.InsertData, offsetBegin, expectRows int) {
|
|
expectInsertData := insertData
|
|
for fieldID, data := range actualInsertData.Data {
|
|
suite.Equal(expectRows, data.RowNum())
|
|
for i := 0; i < expectRows; i++ {
|
|
expect := expectInsertData.Data[fieldID].GetRow(i + offsetBegin)
|
|
actual := data.GetRow(i)
|
|
suite.Equal(expect, actual)
|
|
}
|
|
}
|
|
}
|
|
|
|
res, err := reader.Read()
|
|
suite.NoError(err)
|
|
checkFn(res, 0, suite.numRows)
|
|
}
|
|
|
|
func (suite *ReaderSuite) runWithDefaultValue(dataType schemapb.DataType, elemType schemapb.DataType, oldFormat bool) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "pk",
|
|
IsPrimaryKey: true,
|
|
DataType: suite.pkDataType,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.MaxLengthKey,
|
|
Value: "128",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "vec",
|
|
DataType: suite.vecDataType,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{
|
|
Key: common.DimKey,
|
|
Value: "8",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
// here always set nullable==true just for test, insertData will store validData only nullable==true
|
|
// jsonBytes Marshal from rows, if expectInsertData is nulls and set default value
|
|
// actualData will be default_value
|
|
fieldSchema, err := testutil.CreateFieldWithDefaultValue(dataType, 102, true)
|
|
suite.NoError(err)
|
|
schema.Fields = append(schema.Fields, fieldSchema)
|
|
|
|
insertData, err := testutil.CreateInsertData(schema, suite.numRows)
|
|
suite.NoError(err)
|
|
|
|
rows, err := testutil.CreateInsertDataRowsForJSON(schema, insertData)
|
|
suite.NoError(err)
|
|
|
|
var jsonBytes []byte
|
|
if oldFormat {
|
|
oldRows := make(map[string]any)
|
|
oldRows["rows"] = rows
|
|
jsonBytes, err = json.Marshal(oldRows)
|
|
suite.NoError(err)
|
|
} else {
|
|
jsonBytes, err = json.Marshal(rows)
|
|
suite.NoError(err)
|
|
}
|
|
|
|
cm := mocks.NewChunkManager(suite.T())
|
|
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
|
|
r := importcommon.NewMockReader(string(jsonBytes))
|
|
return r, nil
|
|
})
|
|
reader, err := NewReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
|
|
suite.NoError(err)
|
|
|
|
checkFn := func(actualInsertData *storage.InsertData, offsetBegin, expectRows int) {
|
|
expectInsertData := insertData
|
|
for fieldID, data := range actualInsertData.Data {
|
|
suite.Equal(expectRows, data.RowNum())
|
|
for i := 0; i < expectRows; i++ {
|
|
expect := expectInsertData.Data[fieldID].GetRow(i + offsetBegin)
|
|
actual := data.GetRow(i)
|
|
if expect == nil {
|
|
expect, err = nullutil.GetDefaultValue(fieldSchema)
|
|
suite.NoError(err)
|
|
}
|
|
suite.Equal(expect, actual)
|
|
}
|
|
}
|
|
}
|
|
res, err := reader.Read()
|
|
suite.NoError(err)
|
|
checkFn(res, 0, suite.numRows)
|
|
}
|
|
|
|
func (suite *ReaderSuite) TestReadScalarFields() {
|
|
elementTypes := []schemapb.DataType{
|
|
schemapb.DataType_Bool,
|
|
schemapb.DataType_Int8,
|
|
schemapb.DataType_Int16,
|
|
schemapb.DataType_Int32,
|
|
schemapb.DataType_Int64,
|
|
schemapb.DataType_Float,
|
|
schemapb.DataType_Double,
|
|
schemapb.DataType_VarChar,
|
|
}
|
|
scalarTypes := append(elementTypes, []schemapb.DataType{schemapb.DataType_JSON, schemapb.DataType_Array}...)
|
|
|
|
for _, dataType := range scalarTypes {
|
|
if dataType == schemapb.DataType_Array {
|
|
for _, elementType := range elementTypes {
|
|
suite.run(dataType, elementType, false)
|
|
suite.run(dataType, elementType, true)
|
|
}
|
|
} else {
|
|
suite.run(dataType, schemapb.DataType_None, false)
|
|
suite.run(dataType, schemapb.DataType_None, true)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (suite *ReaderSuite) TestReadScalarFieldsWithDefaultValue() {
|
|
scalarTypes := []schemapb.DataType{
|
|
schemapb.DataType_Bool,
|
|
schemapb.DataType_Int8,
|
|
schemapb.DataType_Int16,
|
|
schemapb.DataType_Int32,
|
|
schemapb.DataType_Int64,
|
|
schemapb.DataType_Float,
|
|
schemapb.DataType_Double,
|
|
schemapb.DataType_VarChar,
|
|
}
|
|
|
|
for _, dataType := range scalarTypes {
|
|
suite.runWithDefaultValue(dataType, schemapb.DataType_None, false)
|
|
suite.runWithDefaultValue(dataType, schemapb.DataType_None, true)
|
|
}
|
|
}
|
|
|
|
func (suite *ReaderSuite) TestStringPK() {
|
|
suite.pkDataType = schemapb.DataType_VarChar
|
|
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, false)
|
|
}
|
|
|
|
func (suite *ReaderSuite) TestVector() {
|
|
suite.vecDataType = schemapb.DataType_BinaryVector
|
|
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, false)
|
|
suite.vecDataType = schemapb.DataType_FloatVector
|
|
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, false)
|
|
suite.vecDataType = schemapb.DataType_Float16Vector
|
|
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, false)
|
|
suite.vecDataType = schemapb.DataType_BFloat16Vector
|
|
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, false)
|
|
suite.vecDataType = schemapb.DataType_SparseFloatVector
|
|
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, false)
|
|
suite.vecDataType = schemapb.DataType_Int8Vector
|
|
suite.run(schemapb.DataType_Int32, schemapb.DataType_None, false)
|
|
}
|
|
|
|
func (suite *ReaderSuite) TestDecodeError() {
|
|
testDecode := func(jsonContent string, isLinesFormat bool, isOldFormat bool, ioErr error, initErr bool, decodeErr bool) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "pk",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
},
|
|
}
|
|
|
|
cm := mocks.NewChunkManager(suite.T())
|
|
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
|
|
r := importcommon.NewMockReader(jsonContent)
|
|
return r, nil
|
|
})
|
|
var reader *reader
|
|
var err error
|
|
if isLinesFormat {
|
|
reader, err = NewLinesReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
|
|
} else {
|
|
reader, err = NewReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
|
|
}
|
|
|
|
if initErr {
|
|
suite.Error(err)
|
|
} else {
|
|
suite.NoError(err)
|
|
}
|
|
|
|
if err == nil {
|
|
reader.isLinesFormat = isLinesFormat
|
|
reader.isOldFormat = isOldFormat
|
|
_, err = reader.Read()
|
|
if decodeErr {
|
|
suite.Error(err)
|
|
} else {
|
|
suite.NoError(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
testOldFormatDecode := func(jsonContent string, ioErr error, initErr bool, decodeErr bool) {
|
|
testDecode(jsonContent, false, true, ioErr, initErr, decodeErr)
|
|
}
|
|
|
|
testListFormatDecode := func(jsonContent string, decodeErr bool) {
|
|
testDecode(jsonContent, false, false, nil, false, decodeErr)
|
|
}
|
|
|
|
testLinesDecode := func(jsonContent string, decodeErr bool) {
|
|
testDecode(jsonContent, true, false, nil, false, decodeErr)
|
|
}
|
|
|
|
// illegal formats
|
|
testOldFormatDecode("", merr.WrapErrImportFailed("error"), true, true)
|
|
testOldFormatDecode("", nil, true, true)
|
|
testOldFormatDecode("a", nil, true, true)
|
|
testOldFormatDecode("2", nil, true, true)
|
|
testOldFormatDecode("{", nil, false, true)
|
|
testOldFormatDecode("{}", nil, false, true) // return EOF
|
|
testOldFormatDecode("{a", nil, false, true)
|
|
testOldFormatDecode("{\"a\":2}", nil, false, true)
|
|
testOldFormatDecode("{\"rows\"a}", nil, false, true)
|
|
testOldFormatDecode("{\"rows\":{}}", nil, false, true)
|
|
testOldFormatDecode("{\"rows\":[a]}", nil, false, true)
|
|
testOldFormatDecode("{\"rows\":[{\"dummy\": 3}]}", nil, false, true)
|
|
testOldFormatDecode("{\"rows\":[{\"pk\": 3}", nil, false, true)
|
|
|
|
// old format
|
|
testOldFormatDecode("{\"rows\":[]}", nil, false, false)
|
|
testOldFormatDecode("{\"rows\":[{\"pk\": 3}]}", nil, false, false)
|
|
testOldFormatDecode("{\n\"rows\":\n[\n{\"pk\": 3},\n{\"pk\": 4}\n]\n}\n", nil, false, false)
|
|
|
|
// new format
|
|
testListFormatDecode("[]", true) // return EOF
|
|
testListFormatDecode("[{\"pk\": 3}]", false)
|
|
testListFormatDecode("[\n{\"pk\": 3},\n{\"pk\": 4}\n]", false)
|
|
|
|
// JSONL/NDJSON format
|
|
testLinesDecode("", true) // return EOF
|
|
testLinesDecode("{\"pk\": 3}", false)
|
|
testLinesDecode("{\"pk\": 3}\n{\"pk\": 4}\n", false)
|
|
}
|
|
|
|
func (suite *ReaderSuite) TestReadCount() {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "pk",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
},
|
|
}
|
|
|
|
rows := make([]any, 0, 100)
|
|
for i := 0; i < 100; i++ {
|
|
row := make(map[string]any)
|
|
row["pk"] = int64(i)
|
|
rows = append(rows, row)
|
|
}
|
|
jsonBytes, err := json.Marshal(rows)
|
|
suite.NoError(err)
|
|
|
|
cm := mocks.NewChunkManager(suite.T())
|
|
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
|
|
r := importcommon.NewMockReader(string(jsonBytes))
|
|
return r, nil
|
|
})
|
|
|
|
// there are 100 rows to be parsed, each read batch is 40 rows
|
|
// buffer size is 320 bytes, each read batch is 40 rows
|
|
// Read() is called for 3 times, 40 + 40 + 20
|
|
reader, err := NewReader(context.Background(), cm, schema, "mockPath", 320)
|
|
suite.NoError(err)
|
|
|
|
data, err := reader.Read()
|
|
suite.NoError(err)
|
|
suite.Equal(40, data.GetRowNum())
|
|
|
|
data, err = reader.Read()
|
|
suite.NoError(err)
|
|
suite.Equal(40, data.GetRowNum())
|
|
|
|
data, err = reader.Read()
|
|
suite.NoError(err)
|
|
suite.Equal(20, data.GetRowNum())
|
|
}
|
|
|
|
func TestJsonReader(t *testing.T) {
|
|
suite.Run(t, new(ReaderSuite))
|
|
}
|
|
|
|
func (suite *ReaderSuite) TestAllowInsertAutoID_KeepUserPK() {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "pk",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
AutoID: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "vec",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: common.DimKey, Value: "8"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
// build rows that explicitly include pk and a valid vec of dim 8
|
|
rows := make([]any, 0, suite.numRows)
|
|
for i := 0; i < suite.numRows; i++ {
|
|
row := make(map[string]any)
|
|
row["pk"] = int64(i + 1)
|
|
vec := make([]float64, 8)
|
|
for j := 0; j < 8; j++ {
|
|
vec[j] = float64(j)
|
|
}
|
|
row["vec"] = vec
|
|
rows = append(rows, row)
|
|
}
|
|
jsonBytes, err := json.Marshal(rows)
|
|
suite.NoError(err)
|
|
|
|
// allow_insert_autoid=false, providing PK should error
|
|
{
|
|
cm := mocks.NewChunkManager(suite.T())
|
|
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
|
|
r := importcommon.NewMockReader(string(jsonBytes))
|
|
return r, nil
|
|
})
|
|
reader, err := NewReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
|
|
suite.NoError(err)
|
|
_, err = reader.Read()
|
|
suite.Error(err)
|
|
suite.Contains(err.Error(), "is auto-generated, no need to provide")
|
|
}
|
|
|
|
// allow_insert_autoid=true, providing PK should be allowed
|
|
{
|
|
schema.Properties = []*commonpb.KeyValuePair{{Key: common.AllowInsertAutoIDKey, Value: "true"}}
|
|
cm := mocks.NewChunkManager(suite.T())
|
|
cm.EXPECT().Reader(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (storage.FileReader, error) {
|
|
r := importcommon.NewMockReader(string(jsonBytes))
|
|
return r, nil
|
|
})
|
|
reader, err := NewReader(context.Background(), cm, schema, "mockPath", math.MaxInt)
|
|
suite.NoError(err)
|
|
_, err = reader.Read()
|
|
suite.NoError(err)
|
|
}
|
|
}
|