From bd88e93bc6fc0187f9b4ffaab15b9d7af9eb53fc Mon Sep 17 00:00:00 2001 From: Yinzuo Jiang Date: Fri, 8 Nov 2024 22:30:55 +0800 Subject: [PATCH] feat: fp32 vector to fp16/bf16 vector conversion for RESTful API RESTful API has 3 handlers. The influenced API are as follows: - Handler. insert - HandlerV1. insert/upsert - HandlerV2. insert/upsert/search issue: #37448 Signed-off-by: Yinzuo Jiang Signed-off-by: Yinzuo Jiang --- .../proxy/httpserver/handler_v2_test.go | 261 +++++++++++++----- .../distributed/proxy/httpserver/utils.go | 126 +++++++-- .../proxy/httpserver/utils_test.go | 108 ++++++-- .../proxy/httpserver/wrap_request.go | 65 ++++- .../proxy/httpserver/wrap_request_test.go | 125 +++++---- internal/json/sonic.go | 2 + .../testcases/test_vector_operations.py | 149 +++++++--- tests/scripts/e2e-restful.sh | 128 ++++++--- .../restful-data/create-collection.json | 29 +- tests/scripts/restful-data/insert-data.json | 22 +- .../{search.json => search-book-intro.json} | 2 +- 11 files changed, 754 insertions(+), 263 deletions(-) rename tests/scripts/restful-data/{search.json => search-book-intro.json} (99%) diff --git a/internal/distributed/proxy/httpserver/handler_v2_test.go b/internal/distributed/proxy/httpserver/handler_v2_test.go index 70c2f63e627d5..d79ec7a1ff968 100644 --- a/internal/distributed/proxy/httpserver/handler_v2_test.go +++ b/internal/distributed/proxy/httpserver/handler_v2_test.go @@ -1577,6 +1577,29 @@ func TestMethodPost(t *testing.T) { } } +func validateTestCases(t *testing.T, testEngine *gin.Engine, queryTestCases []requestBodyTestCase, allowInt64 bool) { + for i, testcase := range queryTestCases { + t.Run(testcase.path, func(t *testing.T) { + bodyReader := bytes.NewReader(testcase.requestBody) + req := httptest.NewRequest(http.MethodPost, versionalV2(EntityCategory, testcase.path), bodyReader) + if allowInt64 { + req.Header.Set(HTTPHeaderAllowInt64, "true") + } + w := httptest.NewRecorder() + testEngine.ServeHTTP(w, req) + assert.Equal(t, http.StatusOK, w.Code, "case %d: ", i, string(testcase.requestBody)) + returnBody := &ReturnErrMsg{} + err := json.Unmarshal(w.Body.Bytes(), returnBody) + assert.Nil(t, err, "case %d: ", i) + assert.Equal(t, testcase.errCode, returnBody.Code, "case %d: ", i, string(testcase.requestBody)) + if testcase.errCode != 0 { + assert.Equal(t, testcase.errMsg, returnBody.Message, "case %d: ", i, string(testcase.requestBody)) + } + fmt.Println(w.Body.String()) + }) + } +} + func TestDML(t *testing.T) { paramtable.Init() // disable rate limit @@ -1692,23 +1715,7 @@ func TestDML(t *testing.T) { requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`), }) - for _, testcase := range queryTestCases { - t.Run(testcase.path, func(t *testing.T) { - bodyReader := bytes.NewReader(testcase.requestBody) - req := httptest.NewRequest(http.MethodPost, versionalV2(EntityCategory, testcase.path), bodyReader) - w := httptest.NewRecorder() - testEngine.ServeHTTP(w, req) - assert.Equal(t, http.StatusOK, w.Code) - returnBody := &ReturnErrMsg{} - err := json.Unmarshal(w.Body.Bytes(), returnBody) - assert.Nil(t, err) - assert.Equal(t, testcase.errCode, returnBody.Code) - if testcase.errCode != 0 { - assert.Equal(t, testcase.errMsg, returnBody.Message) - } - fmt.Println(w.Body.String()) - }) - } + validateTestCases(t, testEngine, queryTestCases, false) } func TestAllowInt64(t *testing.T) { @@ -1736,24 +1743,164 @@ func TestAllowInt64(t *testing.T) { mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Once() mp.EXPECT().Upsert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, UpsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Once() - for _, testcase := range queryTestCases { - t.Run(testcase.path, func(t *testing.T) { - bodyReader := bytes.NewReader(testcase.requestBody) - req := httptest.NewRequest(http.MethodPost, versionalV2(EntityCategory, testcase.path), bodyReader) - req.Header.Set(HTTPHeaderAllowInt64, "true") - w := httptest.NewRecorder() - testEngine.ServeHTTP(w, req) - assert.Equal(t, http.StatusOK, w.Code) - returnBody := &ReturnErrMsg{} - err := json.Unmarshal(w.Body.Bytes(), returnBody) - assert.Nil(t, err) - assert.Equal(t, testcase.errCode, returnBody.Code) - if testcase.errCode != 0 { - assert.Equal(t, testcase.errMsg, returnBody.Message) - } - fmt.Println(w.Body.String()) - }) + validateTestCases(t, testEngine, queryTestCases, true) +} + +func generateCollectionSchemaWithVectorFields() *schemapb.CollectionSchema { + collSchema := generateCollectionSchema(schemapb.DataType_Int64, false) + binaryVectorField := generateVectorFieldSchema(schemapb.DataType_BinaryVector) + binaryVectorField.Name = "binaryVector" + float16VectorField := generateVectorFieldSchema(schemapb.DataType_Float16Vector) + float16VectorField.Name = "float16Vector" + bfloat16VectorField := generateVectorFieldSchema(schemapb.DataType_BFloat16Vector) + bfloat16VectorField.Name = "bfloat16Vector" + sparseFloatVectorField := generateVectorFieldSchema(schemapb.DataType_SparseFloatVector) + sparseFloatVectorField.Name = "sparseFloatVector" + collSchema.Fields = append(collSchema.Fields, binaryVectorField) + collSchema.Fields = append(collSchema.Fields, float16VectorField) + collSchema.Fields = append(collSchema.Fields, bfloat16VectorField) + collSchema.Fields = append(collSchema.Fields, sparseFloatVectorField) + return collSchema +} + +func TestFp16Bf16Vectors(t *testing.T) { + paramtable.Init() + // disable rate limit + paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false") + defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key) + mp := mocks.NewMockProxy(t) + collSchema := generateCollectionSchemaWithVectorFields() + testEngine := initHTTPServerV2(mp, false) + queryTestCases := []requestBodyTestCase{} + for _, path := range []string{InsertAction, UpsertAction} { + queryTestCases = append(queryTestCases, + requestBodyTestCase{ + path: path, + requestBody: []byte( + `{ + "collectionName": "book", + "data": [ + { + "book_id": 0, + "word_count": 0, + "book_intro": [0.11825, 0.6], + "binaryVector": "AQ==", + "float16Vector": [3.0], + "bfloat16Vector": [4.4, 442], + "sparseFloatVector": {"1": 0.1, "2": 0.44} + } + ] + }`), + errCode: 1804, + errMsg: "fail to deal the insert data, error: []byte size 2 doesn't equal to vector dimension 2 of Float16Vector", + }, requestBodyTestCase{ + path: path, + requestBody: []byte( + `{ + "collectionName": "book", + "data": [ + { + "book_id": 0, + "word_count": 0, + "book_intro": [0.11825, 0.6], + "binaryVector": "AQ==", + "float16Vector": [3, 3.0], + "bfloat16Vector": [4.4, 442], + "sparseFloatVector": {"1": 0.1, "2": 0.44} + } + ] + }`), + }, requestBodyTestCase{ + path: path, + // [3, 3] shouble be converted to [float(3), float(3)] + requestBody: []byte( + `{ + "collectionName": "book", + "data": [ + { + "book_id": 0, + "word_count": 0, + "book_intro": [0.11825, 0.6], + "binaryVector": "AQ==", + "float16Vector": [3, 3], + "bfloat16Vector": [4.4, 442], + "sparseFloatVector": {"1": 0.1, "2": 0.44} + } + ] + }`), + }, requestBodyTestCase{ + path: path, + requestBody: []byte( + `{ + "collectionName": "book", + "data": [ + { + "book_id": 0, + "word_count": 0, + "book_intro": [0.11825, 0.6], + "binaryVector": "AQ==", + "float16Vector": "AQIDBA==", + "bfloat16Vector": "AQIDBA==", + "sparseFloatVector": {"1": 0.1, "2": 0.44} + } + ] + }`), + }, requestBodyTestCase{ + path: path, + requestBody: []byte( + `{ + "collectionName": "book", + "data": [ + { + "book_id": 0, + "word_count": 0, + "book_intro": [0.11825, 0.6], + "binaryVector": "AQ==", + "float16Vector": [3, 3.0, 3], + "bfloat16Vector": [4.4, 442, 44], + "sparseFloatVector": {"1": 0.1, "2": 0.44} + } + ] + }`), + errMsg: "fail to deal the insert data, error: []byte size 6 doesn't equal to vector dimension 2 of Float16Vector", + errCode: 1804, + }, requestBodyTestCase{ + path: path, + requestBody: []byte( + `{ + "collectionName": "book", + "data": [ + { + "book_id": 0, + "word_count": 0, + "book_intro": [0.11825, 0.6], + "binaryVector": "AQ==", + "float16Vector": "AQIDBA==", + "bfloat16Vector": [4.4, 442], + "sparseFloatVector": {"1": 0.1, "2": 0.44} + }, + { + "book_id": 1, + "word_count": 0, + "book_intro": [0.11825, 0.6], + "binaryVector": "AQ==", + "float16Vector": [3.1, 3.1], + "bfloat16Vector": "AQIDBA==", + "sparseFloatVector": {"3": 1.1, "2": 0.44} + } + ] + }`), + }) } + mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{ + CollectionName: DefaultCollectionName, + Schema: collSchema, + ShardsNum: ShardNumDefault, + Status: &StatusSuccess, + }, nil).Times(len(queryTestCases)) + mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Times(4) + mp.EXPECT().Upsert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Times(4) + validateTestCases(t, testEngine, queryTestCases, false) } func TestSearchV2(t *testing.T) { @@ -1788,26 +1935,14 @@ func TestSearchV2(t *testing.T) { Ids: generateIDs(schemapb.DataType_Int64, 3), Scores: DefaultScores, }}, nil).Once() - mp.EXPECT().HybridSearch(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{TopK: int64(0)}}, nil).Times(3) - collSchema := generateCollectionSchema(schemapb.DataType_Int64, false) - binaryVectorField := generateVectorFieldSchema(schemapb.DataType_BinaryVector) - binaryVectorField.Name = "binaryVector" - float16VectorField := generateVectorFieldSchema(schemapb.DataType_Float16Vector) - float16VectorField.Name = "float16Vector" - bfloat16VectorField := generateVectorFieldSchema(schemapb.DataType_BFloat16Vector) - bfloat16VectorField.Name = "bfloat16Vector" - sparseFloatVectorField := generateVectorFieldSchema(schemapb.DataType_SparseFloatVector) - sparseFloatVectorField.Name = "sparseFloatVector" - collSchema.Fields = append(collSchema.Fields, binaryVectorField) - collSchema.Fields = append(collSchema.Fields, float16VectorField) - collSchema.Fields = append(collSchema.Fields, bfloat16VectorField) - collSchema.Fields = append(collSchema.Fields, sparseFloatVectorField) + mp.EXPECT().HybridSearch(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{TopK: int64(0)}}, nil).Times(4) + collSchema := generateCollectionSchemaWithVectorFields() mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{ CollectionName: DefaultCollectionName, Schema: collSchema, ShardsNum: ShardNumDefault, Status: &StatusSuccess, - }, nil).Times(10) + }, nil).Times(11) mp.EXPECT().Search(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{TopK: int64(0)}}, nil).Times(3) mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{ Status: &commonpb.Status{ @@ -1906,6 +2041,15 @@ func TestSearchV2(t *testing.T) { `{"data": ["AQIDBA=="], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` + `], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`), }) + queryTestCases = append(queryTestCases, requestBodyTestCase{ + path: AdvancedSearchAction, + requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` + + `{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3},` + + `{"data": ["AQ=="], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` + + `{"data": [[0.1, 0.23]], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` + + `{"data": [[0.1, 0.43]], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` + + `], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`), + }) queryTestCases = append(queryTestCases, requestBodyTestCase{ path: AdvancedSearchAction, requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` + @@ -1983,22 +2127,5 @@ func TestSearchV2(t *testing.T) { errMsg: "mock", errCode: 1100, // ErrParameterInvalid }) - - for _, testcase := range queryTestCases { - t.Run(testcase.path, func(t *testing.T) { - bodyReader := bytes.NewReader(testcase.requestBody) - req := httptest.NewRequest(http.MethodPost, versionalV2(EntityCategory, testcase.path), bodyReader) - w := httptest.NewRecorder() - testEngine.ServeHTTP(w, req) - assert.Equal(t, http.StatusOK, w.Code) - returnBody := &ReturnErrMsg{} - err := json.Unmarshal(w.Body.Bytes(), returnBody) - assert.Nil(t, err) - assert.Equal(t, testcase.errCode, returnBody.Code) - if testcase.errCode != 0 { - assert.Equal(t, testcase.errMsg, returnBody.Message) - } - fmt.Println(w.Body.String()) - }) - } + validateTestCases(t, testEngine, queryTestCases, false) } diff --git a/internal/distributed/proxy/httpserver/utils.go b/internal/distributed/proxy/httpserver/utils.go index 17facf5d0f3a7..a7b515cc48ab5 100644 --- a/internal/distributed/proxy/httpserver/utils.go +++ b/internal/distributed/proxy/httpserver/utils.go @@ -324,27 +324,25 @@ func checkAndSetData(body string, collSchema *schemapb.CollectionSchema) (error, } reallyData[fieldName] = sparseVec case schemapb.DataType_Float16Vector: - if dataString == "" { - return merr.WrapErrParameterInvalid(schemapb.DataType_name[int32(fieldType)], "", "missing vector field: "+fieldName), reallyDataArray, validDataMap - } - vectorStr := gjson.Get(data.Raw, fieldName).Raw - var vectorArray []byte - err := json.Unmarshal([]byte(vectorStr), &vectorArray) - if err != nil { - return merr.WrapErrParameterInvalid(schemapb.DataType_name[int32(fieldType)], dataString, err.Error()), reallyDataArray, validDataMap - } - reallyData[fieldName] = vectorArray + fallthrough case schemapb.DataType_BFloat16Vector: if dataString == "" { return merr.WrapErrParameterInvalid(schemapb.DataType_name[int32(fieldType)], "", "missing vector field: "+fieldName), reallyDataArray, validDataMap } vectorStr := gjson.Get(data.Raw, fieldName).Raw - var vectorArray []byte - err := json.Unmarshal([]byte(vectorStr), &vectorArray) - if err != nil { - return merr.WrapErrParameterInvalid(schemapb.DataType_name[int32(fieldType)], dataString, err.Error()), reallyDataArray, validDataMap + // unmarshal []float32 first to make sure `[3, 3]` is []float instead of []byte + var float32Array []float32 + err := json.Unmarshal([]byte(vectorStr), &float32Array) + if err == nil { + reallyData[fieldName] = float32Array + } else { + var vectorArray []byte + err = json.Unmarshal([]byte(vectorStr), &vectorArray) + if err != nil { + return merr.WrapErrParameterInvalid(schemapb.DataType_name[int32(fieldType)], dataString, err.Error()), reallyDataArray, validDataMap + } + reallyData[fieldName] = vectorArray } - reallyData[fieldName] = vectorArray case schemapb.DataType_Bool: result, err := cast.ToBoolE(dataString) if err != nil { @@ -584,24 +582,46 @@ func convertFloatVectorToArray(vector [][]float32, dim int64) ([]float32, error) } func convertBinaryVectorToArray(vector [][]byte, dim int64, dataType schemapb.DataType) ([]byte, error) { - binaryArray := make([]byte, 0) var bytesLen int64 + var float32BytesLen int64 = -1 switch dataType { case schemapb.DataType_BinaryVector: bytesLen = dim / 8 case schemapb.DataType_Float16Vector: bytesLen = dim * 2 + float32BytesLen = dim * 4 case schemapb.DataType_BFloat16Vector: bytesLen = dim * 2 + float32BytesLen = dim * 4 } + binaryArray := make([]byte, 0, len(vector)*int(bytesLen)) for _, arr := range vector { - if int64(len(arr)) != bytesLen { + switch int64(len(arr)) { + case float32BytesLen: + switch dataType { + // convert float32 to float16 + case schemapb.DataType_Float16Vector: + for i := int64(0); i < float32BytesLen; i += 4 { + f32 := typeutil.BytesToFloat32(arr[i : i+4]) + f16Bytes := typeutil.Float32ToFloat16Bytes(f32) + binaryArray = append(binaryArray, f16Bytes...) + } + // convert float32 to bfloat16 + case schemapb.DataType_BFloat16Vector: + for i := int64(0); i < float32BytesLen; i += 4 { + f32 := typeutil.BytesToFloat32(arr[i : i+4]) + f16Bytes := typeutil.Float32ToBFloat16Bytes(f32) + binaryArray = append(binaryArray, f16Bytes...) + } + } + case bytesLen: + for i := int64(0); i < bytesLen; i++ { + binaryArray = append(binaryArray, arr[i]) + } + default: return nil, fmt.Errorf("[]byte size %d doesn't equal to vector dimension %d of %s", len(arr), dim, schemapb.DataType_name[int32(dataType)]) } - for i := int64(0); i < bytesLen; i++ { - binaryArray = append(binaryArray, arr[i]) - } } return binaryArray, nil } @@ -785,9 +805,25 @@ func anyToColumns(rows []map[string]interface{}, validDataMap map[string][]bool, case schemapb.DataType_BinaryVector: nameColumns[field.Name] = append(nameColumns[field.Name].([][]byte), candi.v.Interface().([]byte)) case schemapb.DataType_Float16Vector: - nameColumns[field.Name] = append(nameColumns[field.Name].([][]byte), candi.v.Interface().([]byte)) + switch candi.v.Interface().(type) { + case []byte: + nameColumns[field.Name] = append(nameColumns[field.Name].([][]byte), candi.v.Interface().([]byte)) + case []float32: + vec := serializeToFloat16(candi.v.Interface().([]float32)) + nameColumns[field.Name] = append(nameColumns[field.Name].([][]byte), vec) + default: + return nil, fmt.Errorf("invalid type(%v) of field(%v) ", field.DataType, field.Name) + } case schemapb.DataType_BFloat16Vector: - nameColumns[field.Name] = append(nameColumns[field.Name].([][]byte), candi.v.Interface().([]byte)) + switch candi.v.Interface().(type) { + case []byte: + nameColumns[field.Name] = append(nameColumns[field.Name].([][]byte), candi.v.Interface().([]byte)) + case []float32: + vec := serializeToBFloat16(candi.v.Interface().([]float32)) + nameColumns[field.Name] = append(nameColumns[field.Name].([][]byte), vec) + default: + return nil, fmt.Errorf("invalid type(%v) of field(%v) ", field.DataType, field.Name) + } case schemapb.DataType_SparseFloatVector: content := candi.v.Interface().([]byte) rowSparseDim := typeutil.SparseFloatRowDim(content) @@ -1024,6 +1060,8 @@ func anyToColumns(rows []map[string]interface{}, validDataMap map[string][]bool, } // --------------------- search param --------------------- // +// serialize serialize vector into byte slice, used in search placeholder +// LittleEndian is used for convention func serialize(fv []float32) []byte { data := make([]byte, 0, 4*len(fv)) // float32 occupies 4 bytes buf := make([]byte, 4) @@ -1034,6 +1072,24 @@ func serialize(fv []float32) []byte { return data } +// serializeToFloat16 converts float32 vector `fv` to float16 vector +func serializeToFloat16(fv []float32) []byte { + data := make([]byte, 0, 2*len(fv)) // float16 occupies 2 bytes + for _, f := range fv { + data = append(data, typeutil.Float32ToFloat16Bytes(f)...) + } + return data +} + +// serializeToBFloat16 converts float32 vector `fv` to bfloat16 vector +func serializeToBFloat16(fv []float32) []byte { + data := make([]byte, 0, 2*len(fv)) // bfloat16 occupies 2 bytes + for _, f := range fv { + data = append(data, typeutil.Float32ToBFloat16Bytes(f)...) + } + return data +} + func serializeFloatVectors(vectors []gjson.Result, dataType schemapb.DataType, dimension, bytesLen int64) ([][]byte, error) { values := make([][]byte, 0) for _, vector := range vectors { @@ -1054,7 +1110,7 @@ func serializeFloatVectors(vectors []gjson.Result, dataType schemapb.DataType, d func serializeByteVectors(vectorStr string, dataType schemapb.DataType, dimension, bytesLen int64) ([][]byte, error) { values := make([][]byte, 0) - err := json.Unmarshal([]byte(vectorStr), &values) // todo check len == dimension * 1/2/2 + err := json.Unmarshal([]byte(vectorStr), &values) if err != nil { return nil, merr.WrapErrParameterInvalid(schemapb.DataType_name[int32(dataType)], vectorStr, err.Error()) } @@ -1067,6 +1123,26 @@ func serializeByteVectors(vectorStr string, dataType schemapb.DataType, dimensio return values, nil } +func serializeFp16OrBf16Vectors(vectorStr string, dataType schemapb.DataType, dimension, bytesLen int64, serializeFunc func([]float32) []byte) ([][]byte, error) { + // try to unmarshal as [][]float32 first to make sure `[[3, 3]]` is [][]float32 instead of [][]byte + fp32Values := make([][]float32, 0) + err := json.Unmarshal([]byte(vectorStr), &fp32Values) + if err == nil { + values := make([][]byte, 0) + for _, vectorArray := range fp32Values { + if int64(len(vectorArray)) != dimension { + vecStr, _ := json.MarshalToString(vectorArray) + return nil, merr.WrapErrParameterInvalid(schemapb.DataType_name[int32(dataType)], vecStr, + fmt.Sprintf("dimension: %d, but length of []float: %d", dimension, len(vectorArray))) + } + vectorBytes := serializeFunc(vectorArray) + values = append(values, vectorBytes) + } + return values, nil + } + return serializeByteVectors(vectorStr, dataType, dimension, bytesLen) +} + func serializeSparseFloatVectors(vectors []gjson.Result, dataType schemapb.DataType) ([][]byte, error) { values := make([][]byte, 0) for _, vector := range vectors { @@ -1093,10 +1169,10 @@ func convertQueries2Placeholder(body string, dataType schemapb.DataType, dimensi values, err = serializeByteVectors(gjson.Get(body, HTTPRequestData).Raw, dataType, dimension, dimension/8) case schemapb.DataType_Float16Vector: valueType = commonpb.PlaceholderType_Float16Vector - values, err = serializeByteVectors(gjson.Get(body, HTTPRequestData).Raw, dataType, dimension, dimension*2) + values, err = serializeFp16OrBf16Vectors(gjson.Get(body, HTTPRequestData).Raw, dataType, dimension, dimension*2, serializeToFloat16) case schemapb.DataType_BFloat16Vector: valueType = commonpb.PlaceholderType_BFloat16Vector - values, err = serializeByteVectors(gjson.Get(body, HTTPRequestData).Raw, dataType, dimension, dimension*2) + values, err = serializeFp16OrBf16Vectors(gjson.Get(body, HTTPRequestData).Raw, dataType, dimension, dimension*2, serializeToBFloat16) case schemapb.DataType_SparseFloatVector: valueType = commonpb.PlaceholderType_SparseFloatVector values, err = serializeSparseFloatVectors(gjson.Get(body, HTTPRequestData).Array(), dataType) diff --git a/internal/distributed/proxy/httpserver/utils_test.go b/internal/distributed/proxy/httpserver/utils_test.go index 11a20cab78f72..cd914c2833a42 100644 --- a/internal/distributed/proxy/httpserver/utils_test.go +++ b/internal/distributed/proxy/httpserver/utils_test.go @@ -701,6 +701,20 @@ func TestInsertWithDefaultValueField(t *testing.T) { func TestSerialize(t *testing.T) { parameters := []float32{0.11111, 0.22222} assert.Equal(t, "\xa4\x8d\xe3=\xa4\x8dc>", string(serialize(parameters))) + + f16vec := serializeToFloat16(parameters) + assert.Equal(t, 4, len(f16vec)) + // \x1c/ is 0.1111, \x1c3 is 0.2222 + assert.Equal(t, "\x1c/\x1c3", string(f16vec)) + assert.Equal(t, "\x1c/", string(typeutil.Float32ToFloat16Bytes(0.11111))) + assert.Equal(t, "\x1c3", string(typeutil.Float32ToFloat16Bytes(0.22222))) + + bf16vec := serializeToBFloat16(parameters) + assert.Equal(t, 4, len(bf16vec)) + assert.Equal(t, "\xe3=c>", string(bf16vec)) + assert.Equal(t, "\xe3=", string(typeutil.Float32ToBFloat16Bytes(0.11111))) + assert.Equal(t, "c>", string(typeutil.Float32ToBFloat16Bytes(0.22222))) + assert.Equal(t, "\n\x10\n\x02$0\x10e\x1a\b\xa4\x8d\xe3=\xa4\x8dc>", string(vectors2PlaceholderGroupBytes([][]float32{parameters}))) // todo requestBody := "{\"data\": [[0.11111, 0.22222]]}" vectors := gjson.Get(requestBody, HTTPRequestData) @@ -1613,31 +1627,50 @@ func TestVector(t *testing.T) { float16Vector := "vector-float16" bfloat16Vector := "vector-bfloat16" sparseFloatVector := "vector-sparse-float" - row1 := map[string]interface{}{ - FieldBookID: int64(1), - floatVector: []float32{0.1, 0.11}, - binaryVector: []byte{1}, - float16Vector: []byte{1, 1, 11, 11}, - bfloat16Vector: []byte{1, 1, 11, 11}, - sparseFloatVector: map[uint32]float32{0: 0.1, 1: 0.11}, - } - row2 := map[string]interface{}{ - FieldBookID: int64(2), - floatVector: []float32{0.2, 0.22}, - binaryVector: []byte{2}, - float16Vector: []byte{2, 2, 22, 22}, - bfloat16Vector: []byte{2, 2, 22, 22}, - sparseFloatVector: map[uint32]float32{1000: 0.3, 200: 0.44}, + testcaseRows := []map[string]interface{}{ + { + FieldBookID: int64(1), + floatVector: []float32{0.1, 0.11}, + binaryVector: []byte{1}, + float16Vector: []byte{1, 1, 11, 11}, + bfloat16Vector: []byte{1, 1, 11, 11}, + sparseFloatVector: map[uint32]float32{0: 0.1, 1: 0.11}, + }, + { + FieldBookID: int64(2), + floatVector: []float32{0.2, 0.22}, + binaryVector: []byte{2}, + float16Vector: []byte{2, 2, 22, 22}, + bfloat16Vector: []byte{2, 2, 22, 22}, + sparseFloatVector: map[uint32]float32{1000: 0.3, 200: 0.44}, + }, + { + FieldBookID: int64(3), + floatVector: []float32{0.3, 0.33}, + binaryVector: []byte{3}, + float16Vector: []byte{3, 3, 33, 33}, + bfloat16Vector: []byte{3, 3, 33, 33}, + sparseFloatVector: map[uint32]float32{987621: 32190.31, 32189: 0.0001}, + }, + { + FieldBookID: int64(4), + floatVector: []float32{0.4, 0.44}, + binaryVector: []byte{4}, + float16Vector: []float32{0.4, 0.44}, + bfloat16Vector: []float32{0.4, 0.44}, + sparseFloatVector: map[uint32]float32{25: 0.1, 1: 0.11}, + }, + { + FieldBookID: int64(5), + floatVector: []float32{-0.4, -0.44}, + binaryVector: []byte{5}, + float16Vector: []int64{99999999, -99999999}, + bfloat16Vector: []int64{99999999, -99999999}, + sparseFloatVector: map[uint32]float32{1121: 0.1, 3: 0.11}, + }, } - row3 := map[string]interface{}{ - FieldBookID: int64(3), - floatVector: []float32{0.3, 0.33}, - binaryVector: []byte{3}, - float16Vector: []byte{3, 3, 33, 33}, - bfloat16Vector: []byte{3, 3, 33, 33}, - sparseFloatVector: map[uint32]float32{987621: 32190.31, 32189: 0.0001}, - } - body, _ := wrapRequestBody([]map[string]interface{}{row1, row2, row3}) + body, err := wrapRequestBody(testcaseRows) + assert.Nil(t, err) primaryField := generatePrimaryField(schemapb.DataType_Int64, false) floatVectorField := generateVectorFieldSchema(schemapb.DataType_FloatVector) floatVectorField.Name = floatVector @@ -1660,10 +1693,27 @@ func TestVector(t *testing.T) { } err, rows, validRows := checkAndSetData(string(body), collectionSchema) assert.Equal(t, nil, err) - for _, row := range rows { + for i, row := range rows { + assert.Equal(t, 2, len(row[floatVector].([]float32))) assert.Equal(t, 1, len(row[binaryVector].([]byte))) - assert.Equal(t, 4, len(row[float16Vector].([]byte))) - assert.Equal(t, 4, len(row[bfloat16Vector].([]byte))) + if fv, ok := testcaseRows[i][float16Vector].([]float32); ok { + assert.Equal(t, 8, len(row[float16Vector].([]byte))) + assert.Equal(t, serialize(fv), row[float16Vector].([]byte)) + } else if _, ok := testcaseRows[i][float16Vector].([]int64); ok { + assert.Equal(t, 8, len(row[float16Vector].([]byte))) + } else { + assert.Equal(t, 4, len(row[float16Vector].([]byte))) + assert.Equal(t, testcaseRows[i][float16Vector].([]byte), row[float16Vector].([]byte)) + } + if fv, ok := testcaseRows[i][bfloat16Vector].([]float32); ok { + assert.Equal(t, 8, len(row[bfloat16Vector].([]byte))) + assert.Equal(t, serialize(fv), row[bfloat16Vector].([]byte)) + } else if _, ok := testcaseRows[i][bfloat16Vector].([]int64); ok { + assert.Equal(t, 8, len(row[bfloat16Vector].([]byte))) + } else { + assert.Equal(t, 4, len(row[bfloat16Vector].([]byte))) + assert.Equal(t, testcaseRows[i][bfloat16Vector].([]byte), row[bfloat16Vector].([]byte)) + } // all test sparse rows have 2 elements, each should be of 8 bytes assert.Equal(t, 16, len(row[sparseFloatVector].([]byte))) } @@ -1674,7 +1724,7 @@ func TestVector(t *testing.T) { assertError := func(field string, value interface{}) { row := make(map[string]interface{}) - for k, v := range row1 { + for k, v := range testcaseRows[0] { row[k] = v } row[field] = value @@ -1683,8 +1733,6 @@ func TestVector(t *testing.T) { assert.Error(t, err) } - assertError(bfloat16Vector, []int64{99999999, -99999999}) - assertError(float16Vector, []int64{99999999, -99999999}) assertError(binaryVector, []int64{99999999, -99999999}) assertError(floatVector, []float64{math.MaxFloat64, 0}) assertError(sparseFloatVector, map[uint32]float32{0: -0.1, 1: 0.11, 2: 0.12}) diff --git a/internal/distributed/proxy/httpserver/wrap_request.go b/internal/distributed/proxy/httpserver/wrap_request.go index 045dbf13d568e..b73b71e6f5277 100644 --- a/internal/distributed/proxy/httpserver/wrap_request.go +++ b/internal/distributed/proxy/httpserver/wrap_request.go @@ -1,10 +1,8 @@ package httpserver import ( - "encoding/binary" "encoding/json" "fmt" - "math" "github.com/cockroachdb/errors" "google.golang.org/protobuf/proto" @@ -77,6 +75,27 @@ type FieldData struct { FieldID int64 `json:"field_id,omitempty"` } +func (f *FieldData) makePbFloat16OrBfloat16Array(raw json.RawMessage, serializeFunc func([]float32) []byte) ([]byte, int64, error) { + wrappedData := [][]float32{} + err := json.Unmarshal(raw, &wrappedData) + if err != nil { + return nil, 0, newFieldDataError(f.FieldName, err) + } + if len(wrappedData) < 1 { + return nil, 0, errors.New("at least one row for insert") + } + array0 := wrappedData[0] + dim := len(array0) + if dim < 1 { + return nil, 0, errors.New("dim must >= 1") + } + data := make([]byte, 0, len(wrappedData)*dim*2) + for _, fp32Array := range wrappedData { + data = append(data, serializeFunc(fp32Array)...) + } + return data, int64(dim), nil +} + // AsSchemapb converts the FieldData to schemapb.FieldData func (f *FieldData) AsSchemapb() (*schemapb.FieldData, error) { // is scarlar @@ -213,6 +232,34 @@ func (f *FieldData) AsSchemapb() (*schemapb.FieldData, error) { }, }, } + case schemapb.DataType_Float16Vector: + // only support float32 conversion right now + data, dim, err := f.makePbFloat16OrBfloat16Array(raw, serializeToFloat16) + if err != nil { + return nil, err + } + ret.Field = &schemapb.FieldData_Vectors{ + Vectors: &schemapb.VectorField{ + Dim: dim, + Data: &schemapb.VectorField_Float16Vector{ + Float16Vector: data, + }, + }, + } + case schemapb.DataType_BFloat16Vector: + // only support float32 conversion right now + data, dim, err := f.makePbFloat16OrBfloat16Array(raw, serializeToBFloat16) + if err != nil { + return nil, err + } + ret.Field = &schemapb.FieldData_Vectors{ + Vectors: &schemapb.VectorField{ + Dim: dim, + Data: &schemapb.VectorField_Bfloat16Vector{ + Bfloat16Vector: data, + }, + }, + } case schemapb.DataType_SparseFloatVector: var wrappedData []map[string]interface{} err := json.Unmarshal(raw, &wrappedData) @@ -309,7 +356,7 @@ func vector2Bytes(vectors [][]float32) []byte { Values: make([][]byte, 0, len(vectors)), } for _, vector := range vectors { - ph.Values = append(ph.Values, serializeVectors(vector)) + ph.Values = append(ph.Values, serialize(vector)) } phg := &commonpb.PlaceholderGroup{ Placeholders: []*commonpb.PlaceholderValue{ @@ -320,18 +367,6 @@ func vector2Bytes(vectors [][]float32) []byte { return ret } -// Serialize serialize vector into byte slice, used in search placeholder -// LittleEndian is used for convention -func serializeVectors(fv []float32) []byte { - data := make([]byte, 0, 4*len(fv)) // float32 occupies 4 bytes - buf := make([]byte, 4) - for _, f := range fv { - binary.LittleEndian.PutUint32(buf, math.Float32bits(f)) - data = append(data, buf...) - } - return data -} - // WrappedCalcDistanceRequest is the RESTful request body for calc distance type WrappedCalcDistanceRequest struct { Base *commonpb.MsgBase `protobuf:"bytes,1,opt,name=base,proto3" json:"base,omitempty"` diff --git a/internal/distributed/proxy/httpserver/wrap_request_test.go b/internal/distributed/proxy/httpserver/wrap_request_test.go index 4d673fb6bd0ff..81430e1c74292 100644 --- a/internal/distributed/proxy/httpserver/wrap_request_test.go +++ b/internal/distributed/proxy/httpserver/wrap_request_test.go @@ -164,61 +164,76 @@ func TestFieldData_AsSchemapb(t *testing.T) { }) // vectors - - t.Run("floatvector_ok", func(t *testing.T) { - fieldData := FieldData{ - Type: schemapb.DataType_FloatVector, - Field: []byte(`[ - [1.1, 2.2, 3.1], - [1.1, 2.2, 3.1], - [1.1, 2.2, 3.1] - ]`), - } - raw, _ := json.Marshal(fieldData) - json.Unmarshal(raw, &fieldData) - _, err := fieldData.AsSchemapb() - assert.NoError(t, err) - }) - t.Run("floatvector_empty_error", func(t *testing.T) { - fieldData := FieldData{ - Type: schemapb.DataType_FloatVector, - Field: []byte(""), - } - raw, _ := json.Marshal(fieldData) - json.Unmarshal(raw, &fieldData) - _, err := fieldData.AsSchemapb() - assert.Error(t, err) - }) - t.Run("floatvector_dim=0_error", func(t *testing.T) { - fieldData := FieldData{ - Type: schemapb.DataType_FloatVector, - Field: []byte(`[]`), - } - raw, _ := json.Marshal(fieldData) - json.Unmarshal(raw, &fieldData) - _, err := fieldData.AsSchemapb() - assert.Error(t, err) - }) - t.Run("floatvector_vectorTypeError_error", func(t *testing.T) { - fieldData := FieldData{ - Type: schemapb.DataType_FloatVector, - Field: []byte(`["1"]`), - } - raw, _ := json.Marshal(fieldData) - json.Unmarshal(raw, &fieldData) - _, err := fieldData.AsSchemapb() - assert.Error(t, err) - }) - t.Run("floatvector_error", func(t *testing.T) { - fieldData := FieldData{ - Type: schemapb.DataType_FloatVector, - Field: []byte(`["a", "b", "c"]`), - } - raw, _ := json.Marshal(fieldData) - json.Unmarshal(raw, &fieldData) - _, err := fieldData.AsSchemapb() - assert.Error(t, err) - }) + testcases := []struct { + name string + dataType schemapb.DataType + }{ + { + "float", schemapb.DataType_FloatVector, + }, + { + "float16", schemapb.DataType_Float16Vector, + }, + { + "bfloat16", schemapb.DataType_BFloat16Vector, + }, + } + for _, tc := range testcases { + t.Run(tc.name+"vector_ok", func(t *testing.T) { + fieldData := FieldData{ + Type: tc.dataType, + Field: []byte(`[ + [1.1, 2.2, 3.1], + [1.1, 2.2, 3.1], + [1.1, 2.2, 3.1] + ]`), + } + raw, _ := json.Marshal(fieldData) + json.Unmarshal(raw, &fieldData) + _, err := fieldData.AsSchemapb() + assert.NoError(t, err) + }) + t.Run(tc.name+"vector_empty_error", func(t *testing.T) { + fieldData := FieldData{ + Type: tc.dataType, + Field: []byte(""), + } + raw, _ := json.Marshal(fieldData) + json.Unmarshal(raw, &fieldData) + _, err := fieldData.AsSchemapb() + assert.Error(t, err) + }) + t.Run(tc.name+"vector_dim=0_error", func(t *testing.T) { + fieldData := FieldData{ + Type: tc.dataType, + Field: []byte(`[]`), + } + raw, _ := json.Marshal(fieldData) + json.Unmarshal(raw, &fieldData) + _, err := fieldData.AsSchemapb() + assert.Error(t, err) + }) + t.Run(tc.name+"vector_vectorTypeError_error", func(t *testing.T) { + fieldData := FieldData{ + Type: tc.dataType, + Field: []byte(`["1"]`), + } + raw, _ := json.Marshal(fieldData) + json.Unmarshal(raw, &fieldData) + _, err := fieldData.AsSchemapb() + assert.Error(t, err) + }) + t.Run(tc.name+"vector_error", func(t *testing.T) { + fieldData := FieldData{ + Type: tc.dataType, + Field: []byte(`["a", "b", "c"]`), + } + raw, _ := json.Marshal(fieldData) + json.Unmarshal(raw, &fieldData) + _, err := fieldData.AsSchemapb() + assert.Error(t, err) + }) + } t.Run("sparsefloatvector_ok_1", func(t *testing.T) { fieldData := FieldData{ diff --git a/internal/json/sonic.go b/internal/json/sonic.go index 816126b163384..61e067c3a4542 100644 --- a/internal/json/sonic.go +++ b/internal/json/sonic.go @@ -24,6 +24,8 @@ import ( var ( json = sonic.ConfigStd + // MarshalToString is exported by gin/json package. + MarshalToString = json.MarshalToString // Marshal is exported by gin/json package. Marshal = json.Marshal // Unmarshal is exported by gin/json package. diff --git a/tests/restful_client_v2/testcases/test_vector_operations.py b/tests/restful_client_v2/testcases/test_vector_operations.py index 39640ba4f1a92..3f158d5e37506 100644 --- a/tests/restful_client_v2/testcases/test_vector_operations.py +++ b/tests/restful_client_v2/testcases/test_vector_operations.py @@ -161,8 +161,10 @@ def test_insert_entities_with_all_scalar_datatype(self, nb, dim, insert_round, a @pytest.mark.parametrize("enable_dynamic_schema", [True]) @pytest.mark.parametrize("nb", [3000]) @pytest.mark.parametrize("dim", [128]) + @pytest.mark.parametrize("pass_fp32_to_fp16_or_bf16", [True, False]) def test_insert_entities_with_all_vector_datatype(self, nb, dim, insert_round, auto_id, - is_partition_key, enable_dynamic_schema): + is_partition_key, enable_dynamic_schema, + pass_fp32_to_fp16_or_bf16): """ Insert a vector with a simple payload """ @@ -210,9 +212,17 @@ def test_insert_entities_with_all_vector_datatype(self, nb, dim, insert_round, a "word_count": i, "book_describe": f"book_{i}", "float_vector": gen_vector(datatype="FloatVector", dim=dim), - "float16_vector": gen_vector(datatype="Float16Vector", dim=dim), - "bfloat16_vector": gen_vector(datatype="BFloat16Vector", dim=dim), - "binary_vector": gen_vector(datatype="BinaryVector", dim=dim) + "float16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="Float16Vector", dim=dim) + ), + "bfloat16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="BFloat16Vector", dim=dim) + ), + "binary_vector": gen_vector(datatype="BinaryVector", dim=dim), } else: tmp = { @@ -221,8 +231,16 @@ def test_insert_entities_with_all_vector_datatype(self, nb, dim, insert_round, a "word_count": i, "book_describe": f"book_{i}", "float_vector": gen_vector(datatype="FloatVector", dim=dim), - "float16_vector": gen_vector(datatype="Float16Vector", dim=dim), - "bfloat16_vector": gen_vector(datatype="BFloat16Vector", dim=dim), + "float16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="Float16Vector", dim=dim) + ), + "bfloat16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="BFloat16Vector", dim=dim) + ), "binary_vector": gen_vector(datatype="BinaryVector", dim=dim) } if enable_dynamic_schema: @@ -253,8 +271,10 @@ def test_insert_entities_with_all_vector_datatype(self, nb, dim, insert_round, a @pytest.mark.parametrize("enable_dynamic_schema", [True]) @pytest.mark.parametrize("nb", [3000]) @pytest.mark.parametrize("dim", [128]) + @pytest.mark.parametrize("pass_fp32_to_fp16_or_bf16", [True, False]) def test_insert_entities_with_all_vector_datatype_0(self, nb, dim, insert_round, auto_id, - is_partition_key, enable_dynamic_schema): + is_partition_key, enable_dynamic_schema, + pass_fp32_to_fp16_or_bf16): """ Insert a vector with a simple payload """ @@ -307,8 +327,16 @@ def test_insert_entities_with_all_vector_datatype_0(self, nb, dim, insert_round, "book_describe": f"book_{i}", "book_vector": gen_vector(datatype="FloatVector", dim=dim), "float_vector": gen_vector(datatype="FloatVector", dim=dim), - "float16_vector": gen_vector(datatype="Float16Vector", dim=dim), - "bfloat16_vector": gen_vector(datatype="BFloat16Vector", dim=dim), + "float16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="Float16Vector", dim=dim) + ), + "bfloat16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="BFloat16Vector", dim=dim) + ), } else: tmp = { @@ -318,8 +346,16 @@ def test_insert_entities_with_all_vector_datatype_0(self, nb, dim, insert_round, "book_describe": f"book_{i}", "book_vector": gen_vector(datatype="FloatVector", dim=dim), "float_vector": gen_vector(datatype="FloatVector", dim=dim), - "float16_vector": gen_vector(datatype="Float16Vector", dim=dim), - "bfloat16_vector": gen_vector(datatype="BFloat16Vector", dim=dim), + "float16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="Float16Vector", dim=dim) + ), + "bfloat16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="BFloat16Vector", dim=dim) + ), } if enable_dynamic_schema: tmp.update({f"dynamic_field_{i}": i}) @@ -349,8 +385,10 @@ def test_insert_entities_with_all_vector_datatype_0(self, nb, dim, insert_round, @pytest.mark.parametrize("enable_dynamic_schema", [True]) @pytest.mark.parametrize("nb", [3000]) @pytest.mark.parametrize("dim", [128]) + @pytest.mark.parametrize("pass_fp32_to_fp16_or_bf16", [True, False]) def test_insert_entities_with_all_vector_datatype_1(self, nb, dim, insert_round, auto_id, - is_partition_key, enable_dynamic_schema): + is_partition_key, enable_dynamic_schema, + pass_fp32_to_fp16_or_bf16): """ Insert a vector with a simple payload """ @@ -399,8 +437,16 @@ def test_insert_entities_with_all_vector_datatype_1(self, nb, dim, insert_round, "word_count": i, "book_describe": f"book_{i}", "float_vector": gen_vector(datatype="FloatVector", dim=dim), - "float16_vector": gen_vector(datatype="Float16Vector", dim=dim), - "bfloat16_vector": gen_vector(datatype="BFloat16Vector", dim=dim), + "float16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="Float16Vector", dim=dim) + ), + "bfloat16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="BFloat16Vector", dim=dim) + ), } else: tmp = { @@ -409,8 +455,16 @@ def test_insert_entities_with_all_vector_datatype_1(self, nb, dim, insert_round, "word_count": i, "book_describe": f"book_{i}", "float_vector": gen_vector(datatype="FloatVector", dim=dim), - "float16_vector": gen_vector(datatype="Float16Vector", dim=dim), - "bfloat16_vector": gen_vector(datatype="BFloat16Vector", dim=dim), + "float16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="Float16Vector", dim=dim) + ), + "bfloat16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="BFloat16Vector", dim=dim) + ), } if enable_dynamic_schema: tmp.update({f"dynamic_field_{i}": i}) @@ -634,7 +688,6 @@ def test_insert_entities_with_all_json_datatype(self, nb, dim, insert_round, aut assert len(rsp['data']) == 50 - @pytest.mark.L0 class TestInsertVectorNegative(TestBase): @@ -937,8 +990,10 @@ class TestSearchVector(TestBase): @pytest.mark.parametrize("enable_dynamic_schema", [True]) @pytest.mark.parametrize("nb", [3000]) @pytest.mark.parametrize("dim", [16]) + @pytest.mark.parametrize("pass_fp32_to_fp16_or_bf16", [True, False]) def test_search_vector_with_all_vector_datatype(self, nb, dim, insert_round, auto_id, - is_partition_key, enable_dynamic_schema): + is_partition_key, enable_dynamic_schema, + pass_fp32_to_fp16_or_bf16): """ Insert a vector with a simple payload """ @@ -986,8 +1041,16 @@ def test_search_vector_with_all_vector_datatype(self, nb, dim, insert_round, aut "word_count": i, "book_describe": f"book_{i}", "float_vector": gen_vector(datatype="FloatVector", dim=dim), - "float16_vector": gen_vector(datatype="Float16Vector", dim=dim), - "bfloat16_vector": gen_vector(datatype="BFloat16Vector", dim=dim), + "float16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="Float16Vector", dim=dim) + ), + "bfloat16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="BFloat16Vector", dim=dim) + ), "binary_vector": gen_vector(datatype="BinaryVector", dim=dim) } else: @@ -997,8 +1060,16 @@ def test_search_vector_with_all_vector_datatype(self, nb, dim, insert_round, aut "word_count": i, "book_describe": f"book_{i}", "float_vector": gen_vector(datatype="FloatVector", dim=dim), - "float16_vector": gen_vector(datatype="Float16Vector", dim=dim), - "bfloat16_vector": gen_vector(datatype="BFloat16Vector", dim=dim), + "float16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="Float16Vector", dim=dim) + ), + "bfloat16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="BFloat16Vector", dim=dim) + ), "binary_vector": gen_vector(datatype="BinaryVector", dim=dim) } if enable_dynamic_schema: @@ -1985,7 +2056,6 @@ def test_search_vector_with_text_match_filter(self, tokenizer): assert token in d[field] - @pytest.mark.L0 class TestSearchVectorNegative(TestBase): @@ -2210,7 +2280,6 @@ def test_advanced_search_vector_with_multi_float32_vector_datatype(self, nb, dim assert len(rsp['data']) == 10 - @pytest.mark.L0 class TestHybridSearchVector(TestBase): @@ -2318,8 +2387,6 @@ def test_hybrid_search_vector_with_multi_float32_vector_datatype(self, nb, dim, assert len(rsp['data']) == 10 - - @pytest.mark.L0 class TestQueryVector(TestBase): @@ -2463,8 +2530,10 @@ def test_query_entities_with_all_scalar_datatype(self, nb, dim, insert_round, au @pytest.mark.parametrize("enable_dynamic_schema", [True]) @pytest.mark.parametrize("nb", [3000]) @pytest.mark.parametrize("dim", [128]) + @pytest.mark.parametrize("pass_fp32_to_fp16_or_bf16", [True, False]) def test_query_entities_with_all_vector_datatype(self, nb, dim, insert_round, auto_id, - is_partition_key, enable_dynamic_schema): + is_partition_key, enable_dynamic_schema, + pass_fp32_to_fp16_or_bf16): """ Insert a vector with a simple payload """ @@ -2512,8 +2581,16 @@ def test_query_entities_with_all_vector_datatype(self, nb, dim, insert_round, au "word_count": i, "book_describe": f"book_{i}", "float_vector": gen_vector(datatype="FloatVector", dim=dim), - "float16_vector": gen_vector(datatype="Float16Vector", dim=dim), - "bfloat16_vector": gen_vector(datatype="BFloat16Vector", dim=dim), + "float16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="Float16Vector", dim=dim) + ), + "bfloat16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="BFloat16Vector", dim=dim) + ), "binary_vector": gen_vector(datatype="BinaryVector", dim=dim) } else: @@ -2523,8 +2600,16 @@ def test_query_entities_with_all_vector_datatype(self, nb, dim, insert_round, au "word_count": i, "book_describe": f"book_{i}", "float_vector": gen_vector(datatype="FloatVector", dim=dim), - "float16_vector": gen_vector(datatype="Float16Vector", dim=dim), - "bfloat16_vector": gen_vector(datatype="BFloat16Vector", dim=dim), + "float16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="Float16Vector", dim=dim) + ), + "bfloat16_vector": ( + gen_vector(datatype="FloatVector", dim=dim) + if pass_fp32_to_fp16_or_bf16 + else gen_vector(datatype="BFloat16Vector", dim=dim) + ), "binary_vector": gen_vector(datatype="BinaryVector", dim=dim) } if enable_dynamic_schema: @@ -2821,8 +2906,6 @@ def test_query_vector_with_text_match_filter(self, tokenizer): assert token in d[field] - - @pytest.mark.L0 class TestQueryVectorNegative(TestBase): diff --git a/tests/scripts/e2e-restful.sh b/tests/scripts/e2e-restful.sh index 5224fd763c9c1..fda700cf620f6 100755 --- a/tests/scripts/e2e-restful.sh +++ b/tests/scripts/e2e-restful.sh @@ -23,93 +23,159 @@ MILVUS_SERVICE_NAME=$(echo "${MILVUS_HELM_RELEASE_NAME}-milvus.${MILVUS_HELM_NAM MILVUS_SERVICE_ADDRESS="${MILVUS_SERVICE_NAME}:9091" # Create a collection -curl -X 'POST' \ +if curl -X 'POST' \ "http://${MILVUS_SERVICE_ADDRESS}/api/v1/collection" \ -H 'accept: application/json' \ -H 'Content-Type: application/json' \ - -d @${DATA_PATH}/create-collection.json + -d @${DATA_PATH}/create-collection.json | grep -q "error_code" ; then + exit 1 +fi # Has collection -curl -X 'GET' \ +if curl -X 'GET' \ "http://${MILVUS_SERVICE_ADDRESS}/api/v1/collection/existence" \ -H 'accept: application/json' \ -H 'Content-Type: application/json' \ -d '{ "collection_name": "book" - }' + }' | grep -q "error_code" ; then + exit 1 +fi # Check collection details -curl -X 'GET' \ +if curl -X 'GET' \ "http://${MILVUS_SERVICE_ADDRESS}/api/v1/collection" \ -H 'accept: application/json' \ -H 'Content-Type: application/json' \ -d '{ "collection_name": "book" - }' + }' | grep -q "error_code" ; then + exit 1 +fi -# Load collection -curl -X 'POST' \ - "http://${MILVUS_SERVICE_ADDRESS}/api/v1/collection/load" \ - -H 'accept: application/json' \ - -H 'Content-Type: application/json' \ - -d '{ - "collection_name": "book" - }' ### Data # Insert Data -curl -X 'POST' \ +if curl -X 'POST' \ "http://${MILVUS_SERVICE_ADDRESS}/api/v1/entities" \ -H 'accept: application/json' \ -H 'Content-Type: application/json' \ - -d @${DATA_PATH}/insert-data.json + -d @${DATA_PATH}/insert-data.json | grep -q "error_code" ; then + exit 1 +fi -# Build Index -curl -X 'POST' \ +# Build Index for book_intro +if curl -X 'POST' \ "http://${MILVUS_SERVICE_ADDRESS}/api/v1/index" \ -H 'accept: application/json' \ -H 'Content-Type: application/json' \ -d '{ "collection_name": "book", "field_name": "book_intro", + "index_name": "book_intro_index", "extra_params":[ {"key": "metric_type", "value": "L2"}, {"key": "index_type", "value": "IVF_FLAT"}, {"key": "params", "value": "{\"nlist\":1024}"} ] - }' + }' | grep -q "error_code" ; then + exit 1 +fi -# KNN Search -curl -X 'POST' \ - "http://${MILVUS_SERVICE_ADDRESS}/api/v1/search" \ +# Build Index for author_intro +if curl -X 'POST' \ + "http://${MILVUS_SERVICE_ADDRESS}/api/v1/index" \ -H 'accept: application/json' \ -H 'Content-Type: application/json' \ - -d @${DATA_PATH}/search.json + -d '{ + "collection_name": "book", + "field_name": "author_intro", + "index_name": "author_intro_index", + "extra_params":[ + {"key": "metric_type", "value": "L2"}, + {"key": "index_type", "value": "IVF_FLAT"}, + {"key": "params", "value": "{\"nlist\":1024}"} + ] + }' | grep -q "error_code" ; then + exit 1 +fi -# Drop Index -curl -X 'DELETE' \ +# Build Index for comment +if curl -X 'POST' \ "http://${MILVUS_SERVICE_ADDRESS}/api/v1/index" \ -H 'accept: application/json' \ -H 'Content-Type: application/json' \ -d '{ "collection_name": "book", - "field_name": "book_intro" - }' + "field_name": "comment", + "index_name": "comment_index", + "extra_params":[ + {"key": "metric_type", "value": "L2"}, + {"key": "index_type", "value": "IVF_FLAT"}, + {"key": "params", "value": "{\"nlist\":1024}"} + ] + }' | grep -q "error_code" ; then + exit 1 +fi + +# Load collection +if curl -X 'POST' \ + "http://${MILVUS_SERVICE_ADDRESS}/api/v1/collection/load" \ + -H 'accept: application/json' \ + -H 'Content-Type: application/json' \ + -d '{ + "collection_name": "book" + }' | grep -q "error_code" ; then + exit 1 +fi + +# KNN Search +# TODO: search fp16/bf16 +for SEARCH_JSON in search-book-intro ; do +if curl -X 'POST' \ + "http://${MILVUS_SERVICE_ADDRESS}/api/v1/search" \ + -H 'accept: application/json' \ + -H 'Content-Type: application/json' \ + -d @${DATA_PATH}/${SEARCH_JSON}.json | grep -q "error_code" ; then + exit 1 +fi +done # Release collection -curl -X 'DELETE' \ +if curl -X 'DELETE' \ "http://${MILVUS_SERVICE_ADDRESS}/api/v1/collection/load" \ -H 'accept: application/json' \ -H 'Content-Type: application/json' \ -d '{ "collection_name": "book" - }' + }' | grep -q "error_code" ; then + exit 1 +fi + +# Drop Index +for FIELD_NAME in book_intro author_intro search_comment ; do +if curl -X 'DELETE' \ + "http://${MILVUS_SERVICE_ADDRESS}/api/v1/index" \ + -H 'accept: application/json' \ + -H 'Content-Type: application/json' \ + -d "{ + \"collection_name\": \"book\", + \"field_name\": \"${FIELD_NAME}\", + \"index_name\": \"${FIELD_NAME}_index\" + }" | grep -q "error_code" ; then + exit 1 +fi +done # Drop collection -curl -X 'DELETE' \ +if curl -X 'DELETE' \ "http://${MILVUS_SERVICE_ADDRESS}/api/v1/collection" \ -H 'accept: application/json' \ -H 'Content-Type: application/json' \ -d '{ "collection_name": "book" - }' + }' | grep -q "error_code" ; then + exit 1 +fi + +echo "e2e-restful.sh success!" diff --git a/tests/scripts/restful-data/create-collection.json b/tests/scripts/restful-data/create-collection.json index 6075ca2d26a42..c51334b07e8d0 100644 --- a/tests/scripts/restful-data/create-collection.json +++ b/tests/scripts/restful-data/create-collection.json @@ -1,7 +1,6 @@ { "collection_name": "book", "consistency_level": 1, - "db_name": "string", "schema": { "autoID": false, "description": "Test book search", @@ -25,8 +24,34 @@ "value": "2" } ] + }, + { + "name": "author_intro", + "description": "embedded vector of author introduction", + "autoID": false, + "data_type": 102, + "is_primary_key": false, + "type_params": [ + { + "key": "dim", + "value": "2" + } + ] + }, + { + "name": "comment", + "description": "embedded vector of comment", + "autoID": false, + "data_type": 103, + "is_primary_key": false, + "type_params": [ + { + "key": "dim", + "value": "2" + } + ] } ], "name": "book" } -} \ No newline at end of file +} diff --git a/tests/scripts/restful-data/insert-data.json b/tests/scripts/restful-data/insert-data.json index fce9191cde43c..fa5a4c8195bb4 100644 --- a/tests/scripts/restful-data/insert-data.json +++ b/tests/scripts/restful-data/insert-data.json @@ -5,16 +5,30 @@ "field_name": "book_id", "type": 5, "field": [ - 1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100, + 1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100 ] }, { "field_name": "book_intro", "type": 101, "field": [ - [1,1],[2,1],[3,1],[4,1],[5,1],[6,1],[7,1],[8,1],[9,1],[10,1],[11,1],[12,1],[13,1],[14,1],[15,1],[16,1],[17,1],[18,1],[19,1],[20,1],[21,1],[22,1],[23,1],[24,1],[25,1],[26,1],[27,1],[28,1],[29,1],[30,1],[31,1],[32,1],[33,1],[34,1],[35,1],[36,1],[37,1],[38,1],[39,1],[40,1],[41,1],[42,1],[43,1],[44,1],[45,1],[46,1],[47,1],[48,1],[49,1],[50,1],[51,1],[52,1],[53,1],[54,1],[55,1],[56,1],[57,1],[58,1],[59,1],[60,1],[61,1],[62,1],[63,1],[64,1],[65,1],[66,1],[67,1],[68,1],[69,1],[70,1],[71,1],[72,1],[73,1],[74,1],[75,1],[76,1],[77,1],[78,1],[79,1],[80,1],[81,1],[82,1],[83,1],[84,1],[85,1],[86,1],[87,1],[88,1],[89,1],[90,1],[91,1],[92,1],[93,1],[94,1],[95,1],[96,1],[97,1],[98,1],[99,1],[100,1], + [1,1],[2,1],[3,1],[4,1],[5,1],[6,1],[7,1],[8,1],[9,1],[10,1],[11,1],[12,1],[13,1],[14,1],[15,1],[16,1],[17,1],[18,1],[19,1],[20,1],[21,1],[22,1],[23,1],[24,1],[25,1],[26,1],[27,1],[28,1],[29,1],[30,1],[31,1],[32,1],[33,1],[34,1],[35,1],[36,1],[37,1],[38,1],[39,1],[40,1],[41,1],[42,1],[43,1],[44,1],[45,1],[46,1],[47,1],[48,1],[49,1],[50,1],[51,1],[52,1],[53,1],[54,1],[55,1],[56,1],[57,1],[58,1],[59,1],[60,1],[61,1],[62,1],[63,1],[64,1],[65,1],[66,1],[67,1],[68,1],[69,1],[70,1],[71,1],[72,1],[73,1],[74,1],[75,1],[76,1],[77,1],[78,1],[79,1],[80,1],[81,1],[82,1],[83,1],[84,1],[85,1],[86,1],[87,1],[88,1],[89,1],[90,1],[91,1],[92,1],[93,1],[94,1],[95,1],[96,1],[97,1],[98,1],[99,1],[100,1] + ] + }, + { + "field_name": "author_intro", + "type": 102, + "field": [ + [1.0,1.0],[2,1.0],[3,1.0],[4,1.0],[5,1.0],[6,1.0],[7,1.0],[8,1.0],[9,1.0],[1.00,1.0],[1.010,1.0],[1.02,1.0],[1.03,1.0],[1.04,1.0],[1.05,1.0],[1.06,1.0],[1.07,1.0],[1.08,1.0],[1.09,1.0],[20,1.0],[21.0,1.0],[22,1.0],[23,1.0],[24,1.0],[25,1.0],[26,1.0],[27,1.0],[28,1.0],[29,1.0],[30,1.0],[31.0,1.0],[32,1.0],[33,1.0],[34,1.0],[35,1.0],[36,1.0],[37,1.0],[38,1.0],[39,1.0],[40,1.0],[41.0,1.0],[42,1.0],[43,1.0],[44,1.0],[45,1.0],[46,1.0],[47,1.0],[48,1.0],[49,1.0],[50,1.0],[51.0,1.0],[52,1.0],[53,1.0],[54,1.0],[55,1.0],[56,1.0],[57,1.0],[58,1.0],[59,1.0],[60,1.0],[61.0,1.0],[62,1.0],[63,1.0],[64,1.0],[65,1.0],[66,1.0],[67,1.0],[68,1.0],[69,1.0],[70,1.0],[71.0,1.0],[72,1.0],[73,1.0],[74,1.0],[75,1.0],[76,1.0],[77,1.0],[78,1.0],[79,1.0],[80,1.0],[81.0,1.0],[82,1.0],[83,1.0],[84,1.0],[85,1.0],[86,1.0],[87,1.0],[88,1.0],[89,1.0],[90,1.0],[91.0,1.0],[92,1.0],[93,1.0],[94,1.0],[95,1.0],[96,1.0],[97,1.0],[98,1.0],[99,1.0],[1.000,1.0] + ] + }, + { + "field_name": "comment", + "type": 103, + "field": [ + [1.0,1.0],[2,1.0],[3,1.0],[4,1.0],[5,1.0],[6,1.0],[7,1.0],[8,1.0],[9,1.0],[1.00,1.0],[1.010,1.0],[1.02,1.0],[1.03,1.0],[1.04,1.0],[1.05,1.0],[1.06,1.0],[1.07,1.0],[1.08,1.0],[1.09,1.0],[20,1.0],[21.0,1.0],[22,1.0],[23,1.0],[24,1.0],[25,1.0],[26,1.0],[27,1.0],[28,1.0],[29,1.0],[30,1.0],[31.0,1.0],[32,1.0],[33,1.0],[34,1.0],[35,1.0],[36,1.0],[37,1.0],[38,1.0],[39,1.0],[40,1.0],[41.0,1.0],[42,1.0],[43,1.0],[44,1.0],[45,1.0],[46,1.0],[47,1.0],[48,1.0],[49,1.0],[50,1.0],[51.0,1.0],[52,1.0],[53,1.0],[54,1.0],[55,1.0],[56,1.0],[57,1.0],[58,1.0],[59,1.0],[60,1.0],[61.0,1.0],[62,1.0],[63,1.0],[64,1.0],[65,1.0],[66,1.0],[67,1.0],[68,1.0],[69,1.0],[70,1.0],[71.0,1.0],[72,1.0],[73,1.0],[74,1.0],[75,1.0],[76,1.0],[77,1.0],[78,1.0],[79,1.0],[80,1.0],[81.0,1.0],[82,1.0],[83,1.0],[84,1.0],[85,1.0],[86,1.0],[87,1.0],[88,1.0],[89,1.0],[90,1.0],[91.0,1.0],[92,1.0],[93,1.0],[94,1.0],[95,1.0],[96,1.0],[97,1.0],[98,1.0],[99,1.0],[1.000,1.0] ] } ], - "num_rows": 1000 -} \ No newline at end of file + "num_rows": 100 +} diff --git a/tests/scripts/restful-data/search.json b/tests/scripts/restful-data/search-book-intro.json similarity index 99% rename from tests/scripts/restful-data/search.json rename to tests/scripts/restful-data/search-book-intro.json index ebcd37ddabac4..3d9500fac9c59 100644 --- a/tests/scripts/restful-data/search.json +++ b/tests/scripts/restful-data/search-book-intro.json @@ -11,4 +11,4 @@ "vectors": [ [10,5] ], "dsl": "", "dsl_type": 1 -} \ No newline at end of file +}