admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-07 01:28:27 +08:00
Code Issues Packages Projects Releases Wiki Activity
milvus/internal/distributed/proxy/httpserver/utils_test.go
cai.zhang 19346fa389
feat: Geospatial Data Type and GIS Function support for milvus (#44547) 
issue: #43427

This pr's main goal is merge #37417 to milvus 2.5 without conflicts.

# Main Goals

1. Create and describe collections with geospatial type
2. Insert geospatial data into the insert binlog
3. Load segments containing geospatial data into memory
4. Enable query and search can display  geospatial data
5. Support using GIS funtions like ST_EQUALS in query
6. Support R-Tree index for geometry type

# Solution

1. **Add Type**: Modify the Milvus core by adding a Geospatial type in
both the C++ and Go code layers, defining the Geospatial data structure
and the corresponding interfaces.
2. **Dependency Libraries**: Introduce necessary geospatial data
processing libraries. In the C++ source code, use Conan package
management to include the GDAL library. In the Go source code, add the
go-geom library to the go.mod file.
3. **Protocol Interface**: Revise the Milvus protocol to provide
mechanisms for Geospatial message serialization and deserialization.
4. **Data Pipeline**: Facilitate interaction between the client and
proxy using the WKT format for geospatial data. The proxy will convert
all data into WKB format for downstream processing, providing column
data interfaces, segment encapsulation, segment loading, payload
writing, and cache block management.
5. **Query Operators**: Implement simple display and support for filter
queries. Initially, focus on filtering based on spatial relationships
for a single column of geospatial literal values, providing parsing and
execution for query expressions.Now only support brutal search
7. **Client Modification**: Enable the client to handle user input for
geospatial data and facilitate end-to-end testing.Check the modification
in pymilvus.

---------

Signed-off-by: Yinwei Li <yinwei.li@zilliz.com>
Signed-off-by: Cai Zhang <cai.zhang@zilliz.com>
Co-authored-by: ZhuXi <150327960+Yinwei-Yu@users.noreply.github.com>
2025-09-28 19:43:05 +08:00

2737 lines
86 KiB
Go
Raw Blame History

// 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 httpserver
import (
"context"
"fmt"
"math"
"strconv"
"strings"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/tidwall/gjson"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/json"
"github.com/milvus-io/milvus/pkg/v2/common"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
const (
FieldWordCount = "word_count"
FieldBookID = "book_id"
FieldBookIntro = "book_intro"
FieldVarchar = "varchar_field"
)
var DefaultScores = []float32{0.01, 0.04, 0.09}
func generatePrimaryField(datatype schemapb.DataType, autoID bool) *schemapb.FieldSchema {
return &schemapb.FieldSchema{
FieldID: common.StartOfUserFieldID,
Name: FieldBookID,
IsPrimaryKey: true,
Description: "",
DataType: datatype,
AutoID: autoID,
}
}
func generateIDs(dataType schemapb.DataType, num int) *schemapb.IDs {
var intArray []int64
if num == 0 {
intArray = []int64{}
} else {
for i := int64(1); i < int64(num+1); i++ {
intArray = append(intArray, i)
}
}
switch dataType {
case schemapb.DataType_Int64:
return &schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: intArray,
},
},
}
case schemapb.DataType_VarChar:
stringArray := formatInt64(intArray)
return &schemapb.IDs{
IdField: &schemapb.IDs_StrId{
StrId: &schemapb.StringArray{
Data: stringArray,
},
},
}
}
return nil
}
func generateVectorFieldSchema(dataType schemapb.DataType) *schemapb.FieldSchema {
dim := "2"
if dataType == schemapb.DataType_BinaryVector {
dim = "8"
}
typeParams := []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: dim,
},
}
if dataType == schemapb.DataType_SparseFloatVector {
typeParams = nil
}
return &schemapb.FieldSchema{
FieldID: common.StartOfUserFieldID + int64(dataType),
IsPrimaryKey: false,
DataType: dataType,
AutoID: false,
TypeParams: typeParams,
}
}
func generateCollectionSchema(primaryDataType schemapb.DataType, autoID bool, isDynamic bool) *schemapb.CollectionSchema {
primaryField := generatePrimaryField(primaryDataType, autoID)
vectorField := generateVectorFieldSchema(schemapb.DataType_FloatVector)
vectorField.Name = FieldBookIntro
fields := []*schemapb.FieldSchema{
primaryField, {
FieldID: common.StartOfUserFieldID + 1,
Name: FieldWordCount,
IsPrimaryKey: false,
Description: "",
DataType: 5,
AutoID: false,
}, vectorField,
}
if isDynamic {
fields = append(fields, &schemapb.FieldSchema{
FieldID: common.StartOfUserFieldID + 2,
Name: "$meta",
IsPrimaryKey: false,
Description: "",
DataType: 23,
AutoID: false,
IsDynamic: true,
})
}
return &schemapb.CollectionSchema{
Name: DefaultCollectionName,
Description: "",
AutoID: autoID,
Fields: fields,
EnableDynamicField: isDynamic,
}
}
func generateDocInDocOutCollectionSchema(primaryDataType schemapb.DataType) *schemapb.CollectionSchema {
primaryField := generatePrimaryField(primaryDataType, false)
vectorField := generateVectorFieldSchema(schemapb.DataType_SparseFloatVector)
vectorField.Name = FieldBookIntro
vectorField.IsFunctionOutput = true
return &schemapb.CollectionSchema{
Name: DefaultCollectionName,
Description: "",
AutoID: false,
Fields: []*schemapb.FieldSchema{
primaryField, {
FieldID: common.StartOfUserFieldID + 1,
Name: FieldWordCount,
IsPrimaryKey: false,
Description: "",
DataType: 5,
AutoID: false,
}, vectorField, {
FieldID: common.StartOfUserFieldID + 2,
Name: FieldVarchar,
IsPrimaryKey: false,
Description: "",
DataType: schemapb.DataType_VarChar,
AutoID: false,
},
},
Functions: []*schemapb.FunctionSchema{
{
Name: "sum",
Type: schemapb.FunctionType_BM25,
InputFieldNames: []string{FieldVarchar},
OutputFieldNames: []string{FieldBookIntro},
},
},
EnableDynamicField: true,
}
}
func generateIndexes() []*milvuspb.IndexDescription {
return []*milvuspb.IndexDescription{
{
IndexName: DefaultIndexName,
IndexID: 442051985533243300,
Params: []*commonpb.KeyValuePair{
{
Key: common.MetricTypeKey,
Value: DefaultMetricType,
},
{
Key: "index_type",
Value: "IVF_FLAT",
},
{
Key: Params,
Value: "{\"nlist\":1024}",
},
},
State: 3,
FieldName: FieldBookIntro,
},
}
}
func generateVectorFieldData(vectorType schemapb.DataType) schemapb.FieldData {
switch vectorType {
case schemapb.DataType_BinaryVector:
return schemapb.FieldData{
Type: schemapb.DataType_BinaryVector,
FieldName: FieldBookIntro,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: 8,
Data: &schemapb.VectorField_BinaryVector{
BinaryVector: []byte{byte(0), byte(1), byte(2)},
},
},
},
IsDynamic: false,
}
case schemapb.DataType_Float16Vector:
return schemapb.FieldData{
Type: schemapb.DataType_Float16Vector,
FieldName: FieldBookIntro,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: 8,
Data: &schemapb.VectorField_Float16Vector{
Float16Vector: []byte{byte(0), byte(0), byte(1), byte(1), byte(2), byte(2)},
},
},
},
IsDynamic: false,
}
case schemapb.DataType_BFloat16Vector:
return schemapb.FieldData{
Type: schemapb.DataType_BFloat16Vector,
FieldName: FieldBookIntro,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: 8,
Data: &schemapb.VectorField_Bfloat16Vector{
Bfloat16Vector: []byte{byte(0), byte(0), byte(1), byte(1), byte(2), byte(2)},
},
},
},
IsDynamic: false,
}
case schemapb.DataType_FloatVector:
return schemapb.FieldData{
Type: schemapb.DataType_FloatVector,
FieldName: FieldBookIntro,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_FloatVector{
FloatVector: &schemapb.FloatArray{
Data: []float32{0.1, 0.11, 0.2, 0.22, 0.3, 0.33},
},
},
},
},
IsDynamic: false,
}
case schemapb.DataType_SparseFloatVector:
contents := make([][]byte, 0, 3)
contents = append(contents, typeutil.CreateSparseFloatRow([]uint32{1, 2, 3}, []float32{0.1, 0.11, 0.2}))
contents = append(contents, typeutil.CreateSparseFloatRow([]uint32{100, 200, 300}, []float32{10.1, 20.11, 30.2}))
contents = append(contents, typeutil.CreateSparseFloatRow([]uint32{1000, 2000, 3000}, []float32{5000.1, 7000.11, 9000.2}))
return schemapb.FieldData{
Type: schemapb.DataType_SparseFloatVector,
FieldName: FieldBookIntro,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: int64(3001),
Data: &schemapb.VectorField_SparseFloatVector{
SparseFloatVector: &schemapb.SparseFloatArray{
Dim: int64(3001),
Contents: contents,
},
},
},
},
IsDynamic: false,
}
case schemapb.DataType_Int8Vector:
return schemapb.FieldData{
Type: schemapb.DataType_Int8Vector,
FieldName: FieldBookIntro,
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: 2,
Data: &schemapb.VectorField_Int8Vector{
Int8Vector: []byte{0x00, 0x1, 0x2, 0x3, 0x4, 0x5},
},
},
},
IsDynamic: false,
}
default:
panic("unsupported vector type")
}
}
func generateFieldData() []*schemapb.FieldData {
fieldData1 := schemapb.FieldData{
Type: schemapb.DataType_Int64,
FieldName: FieldBookID,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{
LongData: &schemapb.LongArray{
Data: []int64{1, 2, 3},
},
},
},
},
IsDynamic: false,
}
fieldData2 := schemapb.FieldData{
Type: schemapb.DataType_Int64,
FieldName: FieldWordCount,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{
LongData: &schemapb.LongArray{
Data: []int64{1000, 2000, 3000},
},
},
},
},
IsDynamic: false,
}
fieldData3 := generateVectorFieldData(schemapb.DataType_FloatVector)
return []*schemapb.FieldData{&fieldData1, &fieldData2, &fieldData3}
}
func wrapRequestBody(data []map[string]interface{}) ([]byte, error) {
body := map[string]interface{}{}
body["data"] = data
return json.Marshal(body)
}
func generateRawRows(dataType schemapb.DataType) []map[string]interface{} {
row1 := map[string]interface{}{
FieldBookID: int64(1),
FieldWordCount: int64(1000),
FieldBookIntro: []float32{0.1, 0.11},
}
row2 := map[string]interface{}{
FieldBookID: int64(2),
FieldWordCount: int64(2000),
FieldBookIntro: []float32{0.2, 0.22},
}
row3 := map[string]interface{}{
FieldBookID: int64(3),
FieldWordCount: int64(3000),
FieldBookIntro: []float32{0.3, 0.33},
}
if dataType == schemapb.DataType_String {
row1[FieldBookID] = "1"
row2[FieldBookID] = "2"
row3[FieldBookID] = "3"
}
return []map[string]interface{}{row1, row2, row3}
}
func generateRequestBody(dataType schemapb.DataType) ([]byte, error) {
return wrapRequestBody(generateRawRows(dataType))
}
func generateRequestBodyWithArray(dataType schemapb.DataType) ([]byte, error) {
rows := generateRawRows(dataType)
for _, result := range rows {
result["array-bool"] = "[true]"
result["array-int8"] = "[0]"
result["array-int16"] = "[0]"
result["array-int32"] = "[0]"
result["array-int64"] = "[0]"
result["array-float"] = "[0.0]"
result["array-double"] = "[0.0]"
result["array-varchar"] = "[\"\"]"
}
return wrapRequestBody(rows)
}
func generateSearchResult(dataType schemapb.DataType) []map[string]interface{} {
rows := generateRawRows(dataType)
for i, row := range rows {
row[DefaultPrimaryFieldName] = row[FieldBookID]
row[HTTPReturnDistance] = DefaultScores[i]
}
return rows
}
func generateQueryResult64(withDistance bool) []map[string]interface{} {
row1 := map[string]interface{}{
FieldBookID: float64(1),
FieldWordCount: float64(1000),
FieldBookIntro: []float64{0.1, 0.11},
}
row2 := map[string]interface{}{
FieldBookID: float64(2),
FieldWordCount: float64(2000),
FieldBookIntro: []float64{0.2, 0.22},
}
row3 := map[string]interface{}{
FieldBookID: float64(3),
FieldWordCount: float64(3000),
FieldBookIntro: []float64{0.3, 0.33},
}
if withDistance {
row1[HTTPReturnDistance] = float64(0.01)
row2[HTTPReturnDistance] = float64(0.04)
row3[HTTPReturnDistance] = float64(0.09)
}
return []map[string]interface{}{row1, row2, row3}
}
func TestPrintCollectionDetails(t *testing.T) {
coll := generateCollectionSchema(schemapb.DataType_Int64, false, true)
indexes := generateIndexes()
assert.Equal(t, []gin.H{
{
HTTPReturnFieldName: FieldBookID,
HTTPReturnFieldType: "Int64",
HTTPReturnFieldPartitionKey: false,
HTTPReturnFieldClusteringKey: false,
HTTPReturnFieldPrimaryKey: true,
HTTPReturnFieldNullable: false,
HTTPReturnFieldAutoID: false,
HTTPReturnDescription: "",
},
{
HTTPReturnFieldName: FieldWordCount,
HTTPReturnFieldType: "Int64",
HTTPReturnFieldPartitionKey: false,
HTTPReturnFieldClusteringKey: false,
HTTPReturnFieldNullable: false,
HTTPReturnFieldPrimaryKey: false,
HTTPReturnFieldAutoID: false,
HTTPReturnDescription: "",
},
{
HTTPReturnFieldName: FieldBookIntro,
HTTPReturnFieldType: "FloatVector(2)",
HTTPReturnFieldPartitionKey: false,
HTTPReturnFieldClusteringKey: false,
HTTPReturnFieldPrimaryKey: false,
HTTPReturnFieldNullable: false,
HTTPReturnFieldAutoID: false,
HTTPReturnDescription: "",
},
}, printFields(coll.Fields))
assert.Equal(t, []gin.H{
{
HTTPReturnFieldName: FieldBookID,
HTTPReturnFieldType: "Int64",
HTTPReturnFieldPartitionKey: false,
HTTPReturnFieldClusteringKey: false,
HTTPReturnFieldPrimaryKey: true,
HTTPReturnFieldNullable: false,
HTTPReturnFieldAutoID: false,
HTTPReturnDescription: "",
HTTPReturnFieldID: int64(100),
},
{
HTTPReturnFieldName: FieldWordCount,
HTTPReturnFieldType: "Int64",
HTTPReturnFieldPartitionKey: false,
HTTPReturnFieldClusteringKey: false,
HTTPReturnFieldPrimaryKey: false,
HTTPReturnFieldNullable: false,
HTTPReturnFieldAutoID: false,
HTTPReturnDescription: "",
HTTPReturnFieldID: int64(101),
},
{
HTTPReturnFieldName: FieldBookIntro,
HTTPReturnFieldType: "FloatVector",
HTTPReturnFieldPartitionKey: false,
HTTPReturnFieldClusteringKey: false,
HTTPReturnFieldNullable: false,
HTTPReturnFieldPrimaryKey: false,
HTTPReturnFieldAutoID: false,
HTTPReturnDescription: "",
HTTPReturnFieldID: int64(201),
Params: []*commonpb.KeyValuePair{
{Key: Dim, Value: "2"},
},
},
}, printFieldsV2(coll.Fields))
assert.Equal(t, []gin.H{
{
HTTPIndexName: DefaultIndexName,
HTTPIndexField: FieldBookIntro,
HTTPReturnIndexMetricType: DefaultMetricType,
},
}, printIndexes(indexes))
assert.Equal(t, DefaultMetricType, getMetricType(indexes[0].Params))
assert.Equal(t, DefaultMetricType, getMetricType(nil))
fields := []*schemapb.FieldSchema{}
for _, field := range newCollectionSchema(coll).Fields {
if field.DataType == schemapb.DataType_VarChar {
fields = append(fields, field)
} else if field.DataType == schemapb.DataType_Array {
fields = append(fields, field)
}
}
assert.Equal(t, []gin.H{
{
HTTPReturnFieldName: "field-varchar",
HTTPReturnFieldType: "VarChar(10)",
HTTPReturnFieldPartitionKey: false,
HTTPReturnFieldClusteringKey: false,
HTTPReturnFieldPrimaryKey: false,
HTTPReturnFieldNullable: false,
HTTPReturnFieldAutoID: false,
HTTPReturnDescription: "",
},
{
HTTPReturnFieldName: "field-array",
HTTPReturnFieldType: "Array",
HTTPReturnFieldPartitionKey: false,
HTTPReturnFieldClusteringKey: false,
HTTPReturnFieldNullable: false,
HTTPReturnFieldPrimaryKey: false,
HTTPReturnFieldAutoID: false,
HTTPReturnDescription: "",
},
}, printFields(fields))
assert.Equal(t, []gin.H{
{
HTTPReturnFieldName: "field-varchar",
HTTPReturnFieldType: "VarChar",
HTTPReturnFieldPartitionKey: false,
HTTPReturnFieldPrimaryKey: false,
HTTPReturnFieldClusteringKey: false,
HTTPReturnFieldAutoID: false,
HTTPReturnFieldNullable: false,
HTTPReturnDescription: "",
HTTPReturnFieldID: int64(0),
Params: []*commonpb.KeyValuePair{
{Key: common.MaxLengthKey, Value: "10"},
},
},
{
HTTPReturnFieldName: "field-array",
HTTPReturnFieldType: "Array",
HTTPReturnFieldPartitionKey: false,
HTTPReturnFieldClusteringKey: false,
HTTPReturnFieldNullable: false,
HTTPReturnFieldPrimaryKey: false,
HTTPReturnFieldAutoID: false,
HTTPReturnDescription: "",
HTTPReturnFieldID: int64(0),
HTTPReturnFieldElementType: "Bool",
},
}, printFieldsV2(fields))
}
func TestPrimaryField(t *testing.T) {
coll := generateCollectionSchema(schemapb.DataType_Int64, false, true)
primaryField := generatePrimaryField(schemapb.DataType_Int64, false)
field, ok := getPrimaryField(coll)
assert.Equal(t, true, ok)
assert.EqualExportedValues(t, primaryField, field)
assert.Equal(t, "1,2,3", joinArray([]int64{1, 2, 3}))
assert.Equal(t, "1,2,3", joinArray([]string{"1", "2", "3"}))
jsonStr := "{\"id\": [1, 2, 3]}"
idStr := gjson.Get(jsonStr, "id")
rangeStr, err := convertRange(primaryField, idStr)
assert.Equal(t, nil, err)
assert.Equal(t, "1,2,3", rangeStr)
filter, err := checkGetPrimaryKey(coll, idStr)
assert.Equal(t, nil, err)
assert.Equal(t, "book_id in [1,2,3]", filter)
primaryField = generatePrimaryField(schemapb.DataType_VarChar, false)
jsonStr = "{\"id\": [\"1\", \"2\", \"3\"]}"
idStr = gjson.Get(jsonStr, "id")
rangeStr, err = convertRange(primaryField, idStr)
assert.Equal(t, nil, err)
assert.Equal(t, `"1","2","3"`, rangeStr)
coll2 := generateCollectionSchema(schemapb.DataType_VarChar, false, true)
filter, err = checkGetPrimaryKey(coll2, idStr)
assert.Equal(t, nil, err)
assert.Equal(t, `book_id in ["1","2","3"]`, filter)
}
func TestAnyToColumns(t *testing.T) {
t.Run("insert with dynamic field", func(t *testing.T) {
body := []byte("{\"data\": {\"id\": 0, \"book_id\": 1, \"book_intro\": [0.1, 0.2], \"word_count\": 2, \"classified\": false, \"databaseID\": null}}")
req := InsertReq{}
coll := generateCollectionSchema(schemapb.DataType_Int64, false, true)
var err error
err, req.Data, _ = checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, int64(0), req.Data[0]["id"])
assert.Equal(t, int64(1), req.Data[0]["book_id"])
assert.Equal(t, int64(2), req.Data[0]["word_count"])
fieldsData, err := anyToColumns(req.Data, nil, coll, true, false)
assert.Equal(t, nil, err)
assert.Equal(t, true, fieldsData[len(fieldsData)-1].IsDynamic)
assert.Equal(t, schemapb.DataType_JSON, fieldsData[len(fieldsData)-1].Type)
assert.Equal(t, "{\"classified\":false,\"id\":0}", string(fieldsData[len(fieldsData)-1].GetScalars().GetJsonData().GetData()[0]))
})
t.Run("upsert with dynamic field", func(t *testing.T) {
body := []byte("{\"data\": {\"id\": 0, \"book_id\": 1, \"book_intro\": [0.1, 0.2], \"word_count\": 2, \"classified\": false, \"databaseID\": null}}")
req := InsertReq{}
coll := generateCollectionSchema(schemapb.DataType_Int64, false, true)
var err error
err, req.Data, _ = checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, int64(0), req.Data[0]["id"])
assert.Equal(t, int64(1), req.Data[0]["book_id"])
assert.Equal(t, int64(2), req.Data[0]["word_count"])
fieldsData, err := anyToColumns(req.Data, nil, coll, false, false)
assert.Equal(t, nil, err)
assert.Equal(t, true, fieldsData[len(fieldsData)-1].IsDynamic)
assert.Equal(t, schemapb.DataType_JSON, fieldsData[len(fieldsData)-1].Type)
assert.Equal(t, "{\"classified\":false,\"id\":0}", string(fieldsData[len(fieldsData)-1].GetScalars().GetJsonData().GetData()[0]))
})
t.Run("insert with dynamic field, but pass pk when autoid==true", func(t *testing.T) {
body := []byte("{\"data\": {\"id\": 0, \"book_id\": 1, \"book_intro\": [0.1, 0.2], \"word_count\": 2, \"classified\": false, \"databaseID\": null}}")
req := InsertReq{}
coll := generateCollectionSchema(schemapb.DataType_Int64, true, true)
var err error
err, req.Data, _ = checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, int64(0), req.Data[0]["id"])
assert.Equal(t, int64(1), req.Data[0]["book_id"])
assert.Equal(t, int64(2), req.Data[0]["word_count"])
_, err = anyToColumns(req.Data, nil, coll, true, false)
assert.Error(t, err)
assert.Equal(t, true, strings.HasPrefix(err.Error(), "no need to pass pk field"))
})
t.Run("pass more field", func(t *testing.T) {
body := []byte("{\"data\": {\"id\": 0, \"book_id\": 1, \"book_intro\": [0.1, 0.2], \"word_count\": 2, \"classified\": false, \"databaseID\": null}}")
coll := generateCollectionSchema(schemapb.DataType_Int64, true, false)
var err error
err, _, _ = checkAndSetData(body, coll, false)
assert.Error(t, err)
assert.Equal(t, true, strings.HasPrefix(err.Error(), "has pass more fiel"))
})
t.Run("insert with autoid==false", func(t *testing.T) {
body := []byte("{\"data\": {\"book_id\": 1, \"book_intro\": [0.1, 0.2], \"word_count\": 2}}")
req := InsertReq{}
coll := generateCollectionSchema(schemapb.DataType_Int64, false, false)
var err error
err, req.Data, _ = checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, int64(1), req.Data[0]["book_id"])
assert.Equal(t, []float32{0.1, 0.2}, req.Data[0]["book_intro"])
assert.Equal(t, int64(2), req.Data[0]["word_count"])
fieldsData, err := anyToColumns(req.Data, nil, coll, true, false)
assert.Equal(t, nil, err)
assert.Equal(t, 3, len(fieldsData))
assert.Equal(t, false, fieldsData[len(fieldsData)-1].IsDynamic)
})
t.Run("insert with autoid==false but has no pk", func(t *testing.T) {
body := []byte("{\"data\": { \"book_intro\": [0.1, 0.2], \"word_count\": 2}}")
coll := generateCollectionSchema(schemapb.DataType_Int64, false, false)
var err error
err, _, _ = checkAndSetData(body, coll, false)
assert.Error(t, err)
assert.Equal(t, true, strings.HasPrefix(err.Error(), "strconv.ParseInt: parsing \"\": invalid syntax"))
})
t.Run("insert with autoid==true", func(t *testing.T) {
body := []byte("{\"data\": { \"book_intro\": [0.1, 0.2], \"word_count\": 2}}")
req := InsertReq{}
coll := generateCollectionSchema(schemapb.DataType_Int64, true, false)
var err error
err, req.Data, _ = checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, []float32{0.1, 0.2}, req.Data[0]["book_intro"])
assert.Equal(t, int64(2), req.Data[0]["word_count"])
fieldsData, err := anyToColumns(req.Data, nil, coll, true, false)
assert.Equal(t, nil, err)
assert.Equal(t, 2, len(fieldsData))
assert.Equal(t, false, fieldsData[len(fieldsData)-1].IsDynamic)
})
t.Run("upsert with autoid==true", func(t *testing.T) {
body := []byte("{\"data\": {\"book_id\": 1, \"book_intro\": [0.1, 0.2], \"word_count\": 2}}")
req := InsertReq{}
coll := generateCollectionSchema(schemapb.DataType_Int64, true, false)
var err error
err, req.Data, _ = checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, int64(1), req.Data[0]["book_id"])
assert.Equal(t, []float32{0.1, 0.2}, req.Data[0]["book_intro"])
assert.Equal(t, int64(2), req.Data[0]["word_count"])
fieldsData, err := anyToColumns(req.Data, nil, coll, false, false)
assert.Equal(t, nil, err)
assert.Equal(t, 3, len(fieldsData))
assert.Equal(t, false, fieldsData[len(fieldsData)-1].IsDynamic)
})
t.Run("upsert with autoid==false", func(t *testing.T) {
body := []byte("{\"data\": {\"book_id\": 1, \"book_intro\": [0.1, 0.2], \"word_count\": 2}}")
req := InsertReq{}
coll := generateCollectionSchema(schemapb.DataType_Int64, true, false)
var err error
err, req.Data, _ = checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, int64(1), req.Data[0]["book_id"])
assert.Equal(t, []float32{0.1, 0.2}, req.Data[0]["book_intro"])
assert.Equal(t, int64(2), req.Data[0]["word_count"])
fieldsData, err := anyToColumns(req.Data, nil, coll, false, false)
assert.Equal(t, nil, err)
assert.Equal(t, 3, len(fieldsData))
assert.Equal(t, false, fieldsData[len(fieldsData)-1].IsDynamic)
})
t.Run("partial update with inconsistent fields should fail", func(t *testing.T) {
// Create a simple schema with two fields: a and b
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "id",
DataType: schemapb.DataType_Int64,
IsPrimaryKey: true,
AutoID: false,
},
{
FieldID: 101,
Name: "a",
DataType: schemapb.DataType_Int64,
},
{
FieldID: 102,
Name: "b",
DataType: schemapb.DataType_Int64,
},
},
EnableDynamicField: false,
}
// Create two rows: first row updates only field 'a', second row updates only field 'b'
rows := []map[string]interface{}{
{
"id": int64(1),
"a": int64(100), // Only field 'a' is provided
},
{
"id": int64(2),
"b": int64(200), // Only field 'b' is provided
},
}
// Test with partial update = true, this should fail
// because different rows are updating different fields
_, err := anyToColumns(rows, nil, schema, false, true)
assert.Error(t, err)
assert.Contains(t, err.Error(), "has length 1, expected 2")
})
t.Run("partial update with consistent missing fields should succeed", func(t *testing.T) {
// Create a simple schema with two fields: a and b
schema := &schemapb.CollectionSchema{
Name: "test_collection",
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "id",
DataType: schemapb.DataType_Int64,
IsPrimaryKey: true,
AutoID: false,
},
{
FieldID: 101,
Name: "a",
DataType: schemapb.DataType_Int64,
},
{
FieldID: 102,
Name: "b",
DataType: schemapb.DataType_Int64,
Nullable: true, // Make field 'b' nullable
},
},
EnableDynamicField: false,
}
// Create two rows: both rows update only field 'a', field 'b' is missing in both
rows := []map[string]interface{}{
{
"id": int64(1),
"a": int64(100), // Only field 'a' is provided
},
{
"id": int64(2),
"a": int64(200), // Only field 'a' is provided
},
}
// Test with partial update = true, this should succeed
// because the same fields are being updated in all rows
fieldsData, err := anyToColumns(rows, nil, schema, false, true)
assert.NoError(t, err)
assert.NotNil(t, fieldsData)
// Should have id and a fields, but not b (since it's not provided and nullable)
fieldNames := make(map[string]bool)
for _, fd := range fieldsData {
fieldNames[fd.FieldName] = true
}
assert.True(t, fieldNames["id"])
assert.True(t, fieldNames["a"])
// Field 'b' should not be present since it wasn't provided in any row
assert.False(t, fieldNames["b"])
})
}
func TestCheckAndSetData(t *testing.T) {
t.Run("invalid field name with dynamic field", func(t *testing.T) {
body := []byte("{\"data\": {\"id\": 0,\"$meta\": 2,\"book_id\": 1, \"book_intro\": [0.1, 0.2], \"word_count\": 2, \"classified\": false, \"databaseID\": null}}")
coll := generateCollectionSchema(schemapb.DataType_Int64, false, true)
var err error
err, _, _ = checkAndSetData(body, coll, false)
assert.Error(t, err)
assert.Equal(t, true, strings.HasPrefix(err.Error(), "use the invalid field name"))
})
t.Run("without vector", func(t *testing.T) {
body := []byte("{\"data\": {}}")
var err error
primaryField := generatePrimaryField(schemapb.DataType_Int64, true)
floatVectorField := generateVectorFieldSchema(schemapb.DataType_FloatVector)
floatVectorField.Name = "floatVector"
binaryVectorField := generateVectorFieldSchema(schemapb.DataType_BinaryVector)
binaryVectorField.Name = "binaryVector"
float16VectorField := generateVectorFieldSchema(schemapb.DataType_Float16Vector)
float16VectorField.Name = "float16Vector"
bfloat16VectorField := generateVectorFieldSchema(schemapb.DataType_BFloat16Vector)
bfloat16VectorField.Name = "bfloat16Vector"
int8VectorField := generateVectorFieldSchema(schemapb.DataType_Int8Vector)
int8VectorField.Name = "int8Vector"
err, _, _ = checkAndSetData(body, &schemapb.CollectionSchema{
Name: DefaultCollectionName,
Fields: []*schemapb.FieldSchema{
primaryField, floatVectorField,
},
EnableDynamicField: true,
}, false)
assert.Error(t, err)
assert.Equal(t, true, strings.HasPrefix(err.Error(), "missing vector field"))
err, _, _ = checkAndSetData(body, &schemapb.CollectionSchema{
Name: DefaultCollectionName,
Fields: []*schemapb.FieldSchema{
primaryField, binaryVectorField,
},
EnableDynamicField: true,
}, false)
assert.Error(t, err)
assert.Equal(t, true, strings.HasPrefix(err.Error(), "missing vector field"))
err, _, _ = checkAndSetData(body, &schemapb.CollectionSchema{
Name: DefaultCollectionName,
Fields: []*schemapb.FieldSchema{
primaryField, float16VectorField,
},
EnableDynamicField: true,
}, false)
assert.Error(t, err)
assert.Equal(t, true, strings.HasPrefix(err.Error(), "missing vector field"))
err, _, _ = checkAndSetData(body, &schemapb.CollectionSchema{
Name: DefaultCollectionName,
Fields: []*schemapb.FieldSchema{
primaryField, bfloat16VectorField,
},
EnableDynamicField: true,
}, false)
assert.Error(t, err)
assert.Equal(t, true, strings.HasPrefix(err.Error(), "missing vector field"))
err, _, _ = checkAndSetData(body, &schemapb.CollectionSchema{
Name: DefaultCollectionName,
Fields: []*schemapb.FieldSchema{
primaryField, int8VectorField,
},
EnableDynamicField: true,
}, false)
assert.Error(t, err)
assert.Equal(t, true, strings.HasPrefix(err.Error(), "missing vector field"))
})
t.Run("with pk when autoID == True when upsert", func(t *testing.T) {
arrayFieldName := "array-int64"
body := []byte("{\"data\": {\"book_id\": 9999999999999999, \"book_intro\": [0.1, 0.2], \"word_count\": 2, \"" + arrayFieldName + "\": [9999999999999999]}}")
coll := generateCollectionSchema(schemapb.DataType_Int64, true, false)
coll.Fields = append(coll.Fields, &schemapb.FieldSchema{
Name: arrayFieldName,
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
})
err, data, validData := checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, 1, len(data))
assert.Equal(t, 0, len(validData))
})
t.Run("without pk when autoID == True when insert", func(t *testing.T) {
arrayFieldName := "array-int64"
body := []byte("{\"data\": {\"book_intro\": [0.1, 0.2], \"word_count\": 2, \"" + arrayFieldName + "\": [9999999999999999]}}")
coll := generateCollectionSchema(schemapb.DataType_Int64, true, false)
coll.Fields = append(coll.Fields, &schemapb.FieldSchema{
Name: arrayFieldName,
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
})
err, data, validData := checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, 1, len(data))
assert.Equal(t, 0, len(validData))
})
t.Run("with pk when autoID == false", func(t *testing.T) {
arrayFieldName := "array-int64"
body := []byte("{\"data\": {\"book_id\": 9999999999999999, \"book_intro\": [0.1, 0.2], \"word_count\": 2, \"" + arrayFieldName + "\": [9999999999999999]}}")
coll := generateCollectionSchema(schemapb.DataType_Int64, false, false)
coll.Fields = append(coll.Fields, &schemapb.FieldSchema{
Name: arrayFieldName,
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
})
err, data, validData := checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, 1, len(data))
assert.Equal(t, 0, len(validData))
})
}
func TestInsertWithInt64(t *testing.T) {
arrayFieldName := "array-int64"
body := []byte("{\"data\": {\"book_id\": 9999999999999999, \"book_intro\": [0.1, 0.2], \"word_count\": 2, \"" + arrayFieldName + "\": [9999999999999999]}}")
coll := generateCollectionSchema(schemapb.DataType_Int64, false, true)
coll.Fields = append(coll.Fields, &schemapb.FieldSchema{
Name: arrayFieldName,
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
})
err, data, validData := checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, 1, len(data))
assert.Equal(t, 0, len(validData))
assert.Equal(t, int64(9999999999999999), data[0][FieldBookID])
arr, _ := data[0][arrayFieldName].(*schemapb.ScalarField)
assert.Equal(t, int64(9999999999999999), arr.GetLongData().GetData()[0])
}
func TestInsertWithNullableField(t *testing.T) {
arrayFieldName := "array-int64"
coll := generateCollectionSchema(schemapb.DataType_Int64, false, true)
coll.Fields = append(coll.Fields, &schemapb.FieldSchema{
Name: arrayFieldName,
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
})
coll.Fields = append(coll.Fields, &schemapb.FieldSchema{
Name: "nullable",
DataType: schemapb.DataType_Int64,
Nullable: true,
})
body := []byte("{\"data\": [{\"book_id\": 9999999999999999, \"\nullable\": null,\"book_intro\": [0.1, 0.2], \"word_count\": 2, \"" + arrayFieldName + "\": [9999999999999999]},{\"book_id\": 1, \"nullable\": 1,\"book_intro\": [0.3, 0.4], \"word_count\": 2, \"" + arrayFieldName + "\": [9999999999999999]}]")
err, data, validData := checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, 2, len(data))
assert.Equal(t, 1, len(validData))
assert.Equal(t, 2, len(validData["nullable"]))
assert.False(t, validData["nullable"][0])
assert.True(t, validData["nullable"][1])
assert.Equal(t, int64(9999999999999999), data[0][FieldBookID])
arr, _ := data[0][arrayFieldName].(*schemapb.ScalarField)
assert.Equal(t, int64(9999999999999999), arr.GetLongData().GetData()[0])
assert.Equal(t, 4, len(data[0]))
assert.Equal(t, 5, len(data[1]))
fieldData, err := anyToColumns(data, validData, coll, true, false)
assert.Equal(t, nil, err)
assert.Equal(t, len(coll.Fields), len(fieldData))
}
func TestInsertWithDefaultValueField(t *testing.T) {
arrayFieldName := "array-int64"
coll := generateCollectionSchema(schemapb.DataType_Int64, false, true)
coll.Fields = append(coll.Fields, &schemapb.FieldSchema{
Name: arrayFieldName,
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
})
coll.Fields = append(coll.Fields, &schemapb.FieldSchema{
Name: "fid",
DataType: schemapb.DataType_Int64,
DefaultValue: &schemapb.ValueField{
Data: &schemapb.ValueField_LongData{
LongData: 10,
},
},
})
body := []byte("{\"data\": [{\"book_id\": 9999999999999999, \"\fid\": null,\"book_intro\": [0.1, 0.2], \"word_count\": 2, \"" + arrayFieldName + "\": [9999999999999999]},{\"book_id\": 1, \"fid\": 1,\"book_intro\": [0.3, 0.4], \"word_count\": 2, \"" + arrayFieldName + "\": [9999999999999999]}]")
err, data, validData := checkAndSetData(body, coll, false)
assert.Equal(t, nil, err)
assert.Equal(t, 2, len(data))
assert.Equal(t, 1, len(validData))
assert.Equal(t, 2, len(validData["fid"]))
assert.False(t, validData["fid"][0])
assert.True(t, validData["fid"][1])
assert.Equal(t, int64(9999999999999999), data[0][FieldBookID])
arr, _ := data[0][arrayFieldName].(*schemapb.ScalarField)
assert.Equal(t, int64(9999999999999999), arr.GetLongData().GetData()[0])
assert.Equal(t, 4, len(data[0]))
assert.Equal(t, 5, len(data[1]))
fieldData, err := anyToColumns(data, validData, coll, true, false)
assert.Equal(t, nil, err)
assert.Equal(t, len(coll.Fields), len(fieldData))
}
func TestSerialize(t *testing.T) {
parameters := []float32{0.11111, 0.22222}
assert.Equal(t, "\n\x10\n\x02$0\x10e\x1a\b\xa4\x8d\xe3=\xa4\x8dc>", string(vectors2PlaceholderGroupBytes([][]float32{parameters}))) // todo
// test serialize fp32 to {fp32, fp16, bf16}
requestBody := "{\"data\": [[0.11111, 0.22222]]}"
vectors := gjson.Get(requestBody, HTTPRequestData)
// fp32 -> fp32
values, err := serializeFloatVectors(vectors.Raw, schemapb.DataType_FloatVector, 2, -1, typeutil.Float32ArrayToBytes)
assert.Nil(t, err)
placeholderValue := &commonpb.PlaceholderValue{
Tag: "$0",
Type: commonpb.PlaceholderType_FloatVector,
Values: values,
}
bytes, err := proto.Marshal(&commonpb.PlaceholderGroup{
Placeholders: []*commonpb.PlaceholderValue{
placeholderValue,
},
})
assert.Nil(t, err)
assert.Equal(t, "\n\x10\n\x02$0\x10e\x1a\b\xa4\x8d\xe3=\xa4\x8dc>", string(bytes)) // todo
// fp32 -> fp16/bf16
for _, testcase := range []struct {
dataType schemapb.DataType
serializeFunc func([]float32) []byte
byteStr string
}{
{schemapb.DataType_Float16Vector, typeutil.Float32ArrayToFloat16Bytes, "\n\f\n\x02$0\x10e\x1a\x04\x1c/\x1c3"},
{schemapb.DataType_BFloat16Vector, typeutil.Float32ArrayToBFloat16Bytes, "\n\f\n\x02$0\x10e\x1a\x04\xe3=c>"},
} {
values, err = serializeFloatOrByteVectors(vectors, testcase.dataType, 2, testcase.serializeFunc)
assert.Nil(t, err)
placeholderValue := &commonpb.PlaceholderValue{
Tag: "$0",
Type: commonpb.PlaceholderType_FloatVector,
Values: values,
}
bytes, err := proto.Marshal(&commonpb.PlaceholderGroup{
Placeholders: []*commonpb.PlaceholderValue{
placeholderValue,
},
})
assert.Nil(t, err)
assert.Equal(t, testcase.byteStr, string(bytes))
}
for _, dataType := range []schemapb.DataType{schemapb.DataType_BinaryVector, schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector} {
request := map[string]interface{}{
HTTPRequestData: []interface{}{
[]byte{1, 2},
},
}
requestBody, _ := json.Marshal(request)
values, err = serializeByteVectors(gjson.Get(string(requestBody), HTTPRequestData).Raw, dataType, -1, 2)
assert.Nil(t, err)
placeholderValue = &commonpb.PlaceholderValue{
Tag: "$0",
Values: values,
}
_, err = proto.Marshal(&commonpb.PlaceholderGroup{
Placeholders: []*commonpb.PlaceholderValue{
placeholderValue,
},
})
assert.Nil(t, err)
}
{
request := map[string]interface{}{
HTTPRequestData: []interface{}{
[]int8{1, 2},
},
}
requestBody, _ := json.Marshal(request)
values, err = serializeInt8Vectors(gjson.Get(string(requestBody), HTTPRequestData).Raw, schemapb.DataType_Int8Vector, 2, typeutil.Int8ArrayToBytes)
assert.Nil(t, err)
placeholderValue = &commonpb.PlaceholderValue{
Tag: "$0",
Values: values,
}
_, err = proto.Marshal(&commonpb.PlaceholderGroup{
Placeholders: []*commonpb.PlaceholderValue{
placeholderValue,
},
})
assert.Nil(t, err)
}
}
func TestConvertQueries2Placeholder(t *testing.T) {
fp16Req := map[string]interface{}{
HTTPRequestData: []interface{}{
typeutil.Float32ArrayToFloat16Bytes([]float32{0.11111, 0.22222, 0, -1}),
typeutil.Float32ArrayToFloat16Bytes([]float32{1, 1, 1, 1}),
},
}
fp16ReqBody, _ := json.Marshal(fp16Req)
const Float16VecJSON = `{"data":["HC8cMwAAALw=","ADwAPAA8ADw="]}`
assert.Equal(t, Float16VecJSON, string(fp16ReqBody))
bf16Req := map[string]interface{}{
HTTPRequestData: []interface{}{
typeutil.Float32ArrayToBFloat16Bytes([]float32{0.11111, 0.22222, 0, -1}),
typeutil.Float32ArrayToBFloat16Bytes([]float32{1, 1, 1, 1}),
},
}
bf16ReqBody, _ := json.Marshal(bf16Req)
const BFloat16VecJSON = `{"data":["4z1jPgAAgL8=","gD+AP4A/gD8="]}`
assert.Equal(t, BFloat16VecJSON, string(bf16ReqBody))
type testCase struct {
requestBody string
dataType schemapb.DataType
dim int64
placehoderValue func() [][]byte
}
testCases := make([]testCase, 0)
for _, dataType := range []schemapb.DataType{schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector, schemapb.DataType_FloatVector} {
// corner case: empty data
testCases = append(testCases, []testCase{
{
"{\"data\": []}",
dataType,
0,
func() [][]byte {
return [][]byte{}
},
}, {
"{\"data\": []}",
dataType,
100,
func() [][]byte {
return [][]byte{}
},
}, {
"{\"data\": [[], []]}",
dataType,
0,
func() [][]byte {
return [][]byte{{}, {}}
},
},
}...)
}
for _, dataType := range []schemapb.DataType{schemapb.DataType_Float16Vector, schemapb.DataType_BFloat16Vector} {
// corner case: empty float16/bfloat16 vector
testCases = append(testCases, []testCase{
{
`"{"data": ["", ""]}"`,
dataType,
0,
func() [][]byte {
return [][]byte{{}, {}}
},
}, {
`"{"data": [""]}"`,
dataType,
0,
func() [][]byte {
return [][]byte{{}}
},
},
}...)
}
testCases = append(testCases, []testCase{
{
"{\"data\": [[0.11111, 0.22222]]}",
schemapb.DataType_FloatVector,
2,
func() [][]byte {
bv := typeutil.Float32ArrayToBytes([]float32{0.11111, 0.22222})
return [][]byte{bv}
},
}, {
"{\"data\": [[0.11111, 0.22222, 0, -1]]}",
schemapb.DataType_Float16Vector,
4,
func() [][]byte {
bv := typeutil.Float32ArrayToFloat16Bytes([]float32{0.11111, 0.22222, 0, -1})
return [][]byte{bv}
},
}, {
"{\"data\": [[0.11111, 0.22222, 0, -1], [1, 1, 1, 1]]}",
schemapb.DataType_Float16Vector,
4,
func() [][]byte {
bv1 := typeutil.Float32ArrayToFloat16Bytes([]float32{0.11111, 0.22222, 0, -1})
bv2 := typeutil.Float32ArrayToFloat16Bytes([]float32{1, 1, 1, 1})
return [][]byte{bv1, bv2}
},
}, {
"{\"data\": [[0.11111, 0.22222, 0, -1], [1, 1, 1, 1]]}",
schemapb.DataType_BFloat16Vector,
4,
func() [][]byte {
bv1 := typeutil.Float32ArrayToBFloat16Bytes([]float32{0.11111, 0.22222, 0, -1})
bv2 := typeutil.Float32ArrayToBFloat16Bytes([]float32{1, 1, 1, 1})
return [][]byte{bv1, bv2}
},
}, {
Float16VecJSON,
schemapb.DataType_Float16Vector,
4,
func() [][]byte {
bv1 := typeutil.Float32ArrayToFloat16Bytes([]float32{0.11111, 0.22222, 0, -1})
bv2 := typeutil.Float32ArrayToFloat16Bytes([]float32{1, 1, 1, 1})
return [][]byte{bv1, bv2}
},
}, {
BFloat16VecJSON,
schemapb.DataType_BFloat16Vector,
4,
func() [][]byte {
bv1 := typeutil.Float32ArrayToBFloat16Bytes([]float32{0.11111, 0.22222, 0, -1})
bv2 := typeutil.Float32ArrayToBFloat16Bytes([]float32{1, 1, 1, 1})
return [][]byte{bv1, bv2}
},
},
}...)
for _, testcase := range testCases {
phv, err := convertQueries2Placeholder(testcase.requestBody, testcase.dataType, testcase.dim)
assert.Nil(t, err)
assert.Equal(t, testcase.placehoderValue(), phv.GetValues(),
fmt.Sprintf("check equal fail, data: %s, type: %s, dim: %d", testcase.requestBody, testcase.dataType, testcase.dim))
}
for _, testcase := range []testCase{
// mismatched Datatype
{
"{\"data\": [[0.11111, 0.22222, 0, -1], [1, 1, 1, 1]]}",
schemapb.DataType_Float16Vector,
4,
func() [][]byte {
bv1 := typeutil.Float32ArrayToBFloat16Bytes([]float32{0.11111, 0.22222, 0, -1})
bv2 := typeutil.Float32ArrayToBFloat16Bytes([]float32{1, 1, 1, 1})
return [][]byte{bv1, bv2}
},
}, {
"{\"data\": [[0.11111, 0.22222, 0, -1], [1, 1, 1, 1]]}",
schemapb.DataType_BFloat16Vector,
4,
func() [][]byte {
bv1 := typeutil.Float32ArrayToFloat16Bytes([]float32{0.11111, 0.22222, 0, -1})
bv2 := typeutil.Float32ArrayToFloat16Bytes([]float32{1, 1, 1, 1})
return [][]byte{bv1, bv2}
},
},
} {
phv, err := convertQueries2Placeholder(testcase.requestBody, testcase.dataType, testcase.dim)
assert.Nil(t, err)
assert.NotEqual(t, testcase.placehoderValue(), phv.GetValues(),
fmt.Sprintf("check not equal fail, data: %s, type: %s, dim: %d", testcase.requestBody, testcase.dataType, testcase.dim))
}
for _, testcase := range []testCase{
// mismatched dimension
{
"{\"data\": [[0.11111, 0.22222, 0, -1], [1, 1, 1, 1]]}",
schemapb.DataType_Float16Vector,
2,
func() [][]byte {
bv1 := typeutil.Float32ArrayToBFloat16Bytes([]float32{0.11111, 0.22222, 0, -1})
bv2 := typeutil.Float32ArrayToBFloat16Bytes([]float32{1, 1, 1, 1})
return [][]byte{bv1, bv2}
},
}, {
"{\"data\": [[0.11111, 0.22222, 0, -1], [1, 1, 1, 1]]}",
schemapb.DataType_BFloat16Vector,
8,
func() [][]byte {
bv1 := typeutil.Float32ArrayToFloat16Bytes([]float32{0.11111, 0.22222, 0, -1})
bv2 := typeutil.Float32ArrayToFloat16Bytes([]float32{1, 1, 1, 1})
return [][]byte{bv1, bv2}
},
}, {
"{\"data\": [[0.11111, 0.22222, 0, -1], [1, 1, 1]]}",
schemapb.DataType_BFloat16Vector,
4,
func() [][]byte {
bv1 := typeutil.Float32ArrayToFloat16Bytes([]float32{0.11111, 0.22222, 0, -1})
bv2 := typeutil.Float32ArrayToFloat16Bytes([]float32{1, 1, 1, 1})
return [][]byte{bv1, bv2}
},
},
} {
_, err := convertQueries2Placeholder(testcase.requestBody, testcase.dataType, testcase.dim)
assert.NotNil(t, err)
}
}
func compareRow64(m1 map[string]interface{}, m2 map[string]interface{}) bool {
for key, value := range m1 {
if key == FieldBookIntro {
arr1 := value.([]interface{})
arr2 := m2[key].([]float64)
if len(arr1) != len(arr2) {
return false
}
for j, element := range arr1 {
if element != arr2[j] {
return false
}
}
} else if value != m2[key] {
return false
}
}
for key, value := range m2 {
if key == FieldBookIntro {
continue
} else if value != m1[key] {
return false
}
}
return true
}
func compareRow(m1 map[string]interface{}, m2 map[string]interface{}) bool {
for key, value := range m1 {
if key == FieldBookIntro {
arr1 := value.([]float32)
arr2 := m2[key].([]float32)
if len(arr1) != len(arr2) {
return false
}
for j, element := range arr1 {
if element != arr2[j] {
return false
}
}
} else if key == "field-json" {
arr1 := value.(string)
arr2 := m2[key].([]byte)
if arr1 != string(arr2) {
return false
}
} else if key == "field-geometry" {
arr1 := value.(string)
arr2 := m2[key].(string)
if arr2 != arr1 {
return false
}
} else if strings.HasPrefix(key, "array-") {
continue
} else if value != m2[key] {
return false
}
}
for key, value := range m2 {
if (key == FieldBookIntro) || (key == "field-json") || (key == "field-geometry") || (key == "field-array") {
continue
} else if strings.HasPrefix(key, "array-") {
continue
} else if value != m1[key] {
return false
}
}
return true
}
type CompareFunc func(map[string]interface{}, map[string]interface{}) bool
func compareRows(row1 []map[string]interface{}, row2 []map[string]interface{}, compareFunc CompareFunc) bool {
if len(row1) != len(row2) {
return false
}
for i, row := range row1 {
if !compareFunc(row, row2[i]) {
return false
}
}
return true
}
func TestBuildQueryResp(t *testing.T) {
outputFields := []string{FieldBookID, FieldWordCount, "author", "date"}
rows, err := buildQueryResp(int64(0), outputFields, generateFieldData(), generateIDs(schemapb.DataType_Int64, 3), DefaultScores, true, nil) // []*schemapb.FieldData{&fieldData1, &fieldData2, &fieldData3}
assert.Equal(t, nil, err)
exceptRows := generateSearchResult(schemapb.DataType_Int64)
assert.Equal(t, true, compareRows(rows, exceptRows, compareRow))
}
func newCollectionSchema(coll *schemapb.CollectionSchema) *schemapb.CollectionSchema {
fieldSchema1 := schemapb.FieldSchema{
Name: "field-bool",
DataType: schemapb.DataType_Bool,
}
coll.Fields = append(coll.Fields, &fieldSchema1)
fieldSchema2 := schemapb.FieldSchema{
Name: "field-int8",
DataType: schemapb.DataType_Int8,
}
coll.Fields = append(coll.Fields, &fieldSchema2)
fieldSchema3 := schemapb.FieldSchema{
Name: "field-int16",
DataType: schemapb.DataType_Int16,
}
coll.Fields = append(coll.Fields, &fieldSchema3)
fieldSchema4 := schemapb.FieldSchema{
Name: "field-int32",
DataType: schemapb.DataType_Int32,
}
coll.Fields = append(coll.Fields, &fieldSchema4)
fieldSchema5 := schemapb.FieldSchema{
Name: "field-float",
DataType: schemapb.DataType_Float,
}
coll.Fields = append(coll.Fields, &fieldSchema5)
fieldSchema6 := schemapb.FieldSchema{
Name: "field-double",
DataType: schemapb.DataType_Double,
}
coll.Fields = append(coll.Fields, &fieldSchema6)
fieldSchema7 := schemapb.FieldSchema{
Name: "field-string",
DataType: schemapb.DataType_String,
}
coll.Fields = append(coll.Fields, &fieldSchema7)
fieldSchema8 := schemapb.FieldSchema{
Name: "field-varchar",
DataType: schemapb.DataType_VarChar,
TypeParams: []*commonpb.KeyValuePair{
{Key: "max_length", Value: "10"},
},
}
coll.Fields = append(coll.Fields, &fieldSchema8)
fieldSchema9 := schemapb.FieldSchema{
Name: "field-json",
DataType: schemapb.DataType_JSON,
IsDynamic: false,
}
coll.Fields = append(coll.Fields, &fieldSchema9)
fieldSchema10 := schemapb.FieldSchema{
Name: "field-array",
DataType: schemapb.DataType_Array,
IsDynamic: false,
ElementType: schemapb.DataType_Bool,
}
coll.Fields = append(coll.Fields, &fieldSchema10)
fieldSchema11 := schemapb.FieldSchema{
Name: "field-geometry",
DataType: schemapb.DataType_Geometry,
IsDynamic: false,
}
coll.Fields = append(coll.Fields, &fieldSchema11)
return coll
}
func withDynamicField(coll *schemapb.CollectionSchema) *schemapb.CollectionSchema {
fieldSchema11 := schemapb.FieldSchema{
Name: "$meta",
DataType: schemapb.DataType_JSON,
IsDynamic: true,
}
coll.Fields = append(coll.Fields, &fieldSchema11)
return coll
}
func withArrayField(coll *schemapb.CollectionSchema) *schemapb.CollectionSchema {
fieldSchema0 := schemapb.FieldSchema{
Name: "array-bool",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Bool,
}
coll.Fields = append(coll.Fields, &fieldSchema0)
fieldSchema1 := schemapb.FieldSchema{
Name: "array-int8",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int8,
}
coll.Fields = append(coll.Fields, &fieldSchema1)
fieldSchema2 := schemapb.FieldSchema{
Name: "array-int16",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int16,
}
coll.Fields = append(coll.Fields, &fieldSchema2)
fieldSchema3 := schemapb.FieldSchema{
Name: "array-int32",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int32,
}
coll.Fields = append(coll.Fields, &fieldSchema3)
fieldSchema4 := schemapb.FieldSchema{
Name: "array-int64",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
}
coll.Fields = append(coll.Fields, &fieldSchema4)
fieldSchema5 := schemapb.FieldSchema{
Name: "array-float",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Float,
}
coll.Fields = append(coll.Fields, &fieldSchema5)
fieldSchema6 := schemapb.FieldSchema{
Name: "array-double",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Double,
}
coll.Fields = append(coll.Fields, &fieldSchema6)
fieldSchema7 := schemapb.FieldSchema{
Name: "array-varchar",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_VarChar,
}
coll.Fields = append(coll.Fields, &fieldSchema7)
return coll
}
func newFieldData(fieldDatas []*schemapb.FieldData, firstFieldType schemapb.DataType) []*schemapb.FieldData {
fieldData1 := schemapb.FieldData{
Type: schemapb.DataType_Bool,
FieldName: "field-bool",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_BoolData{
BoolData: &schemapb.BoolArray{
Data: []bool{true, true, true},
},
},
},
},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData1)
fieldData2 := schemapb.FieldData{
Type: schemapb.DataType_Int8,
FieldName: "field-int8",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{
IntData: &schemapb.IntArray{
Data: []int32{0, 1, 2},
},
},
},
},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData2)
fieldData3 := schemapb.FieldData{
Type: schemapb.DataType_Int16,
FieldName: "field-int16",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{
IntData: &schemapb.IntArray{
Data: []int32{0, 1, 2},
},
},
},
},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData3)
fieldData4 := schemapb.FieldData{
Type: schemapb.DataType_Int32,
FieldName: "field-int32",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{
IntData: &schemapb.IntArray{
Data: []int32{0, 1, 2},
},
},
},
},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData4)
fieldData5 := schemapb.FieldData{
Type: schemapb.DataType_Float,
FieldName: "field-float",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_FloatData{
FloatData: &schemapb.FloatArray{
Data: []float32{0, 1, 2},
},
},
},
},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData5)
fieldData6 := schemapb.FieldData{
Type: schemapb.DataType_Double,
FieldName: "field-double",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_DoubleData{
DoubleData: &schemapb.DoubleArray{
Data: []float64{0, 1, 2},
},
},
},
},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData6)
fieldData7 := schemapb.FieldData{
Type: schemapb.DataType_String,
FieldName: "field-string",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{"0", "1", "2"},
},
},
},
},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData7)
fieldData8 := schemapb.FieldData{
Type: schemapb.DataType_VarChar,
FieldName: "field-varchar",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{"0", "1", "2"},
},
},
},
},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData8)
fieldData9 := schemapb.FieldData{
Type: schemapb.DataType_JSON,
FieldName: "field-json",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_JsonData{
JsonData: &schemapb.JSONArray{
Data: [][]byte{[]byte(`{"XXX": 0}`), []byte(`{"XXX": 0}`), []byte(`{"XXX": 0}`)},
},
},
},
},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData9)
fieldData10 := schemapb.FieldData{
Type: schemapb.DataType_Array,
FieldName: "field-array",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_ArrayData{
ArrayData: &schemapb.ArrayArray{
Data: []*schemapb.ScalarField{
{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true}}}},
{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true}}}},
{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true}}}},
},
},
},
},
},
IsDynamic: false,
}
fieldData11 := schemapb.FieldData{
Type: schemapb.DataType_JSON,
FieldName: "$meta",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_JsonData{
JsonData: &schemapb.JSONArray{
Data: [][]byte{[]byte(`{"XXX": 0, "YYY": "0"}`), []byte(`{"XXX": 1, "YYY": "1"}`), []byte(`{"XXX": 2, "YYY": "2"}`)},
},
},
},
},
IsDynamic: true,
}
fieldDatas = append(fieldDatas, &fieldData11)
fieldData12 := schemapb.FieldData{
Type: schemapb.DataType_Geometry,
FieldName: "field-geometry",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_GeometryWktData{
GeometryWktData: &schemapb.GeometryWktArray{
Data: []string{
`POINT (30.123 -10.456)`,
`POINT (30.123 -10.456)`,
`POINT (30.123 -10.456)`,
// wkb:{0x01, 0x01, 0x00, 0x00, 0x00, 0xD2, 0x4A, 0x4D, 0x6A, 0x8B, 0x3C, 0x5C, 0x0A, 0x0D, 0x1B, 0x4F, 0x4F, 0x9A, 0x3D, 0x4},
},
},
},
},
},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData12)
switch firstFieldType {
case schemapb.DataType_None:
return fieldDatas
case schemapb.DataType_Bool:
return []*schemapb.FieldData{&fieldData1}
case schemapb.DataType_Int8:
return []*schemapb.FieldData{&fieldData2}
case schemapb.DataType_Int16:
return []*schemapb.FieldData{&fieldData3}
case schemapb.DataType_Int32:
return []*schemapb.FieldData{&fieldData4}
case schemapb.DataType_Float:
return []*schemapb.FieldData{&fieldData5}
case schemapb.DataType_Double:
return []*schemapb.FieldData{&fieldData6}
case schemapb.DataType_String:
return []*schemapb.FieldData{&fieldData7}
case schemapb.DataType_VarChar:
return []*schemapb.FieldData{&fieldData8}
case schemapb.DataType_BinaryVector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_FloatVector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_Float16Vector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_BFloat16Vector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_Int8Vector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_Array:
return []*schemapb.FieldData{&fieldData10}
case schemapb.DataType_JSON:
return []*schemapb.FieldData{&fieldData9}
case schemapb.DataType_Geometry:
return []*schemapb.FieldData{&fieldData12}
case schemapb.DataType_SparseFloatVector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
default:
return []*schemapb.FieldData{
{
FieldName: "wrong-field-type",
Type: firstFieldType,
},
}
}
}
func newNullableFieldData(fieldDatas []*schemapb.FieldData, firstFieldType schemapb.DataType) []*schemapb.FieldData {
fieldData1 := schemapb.FieldData{
Type: schemapb.DataType_Bool,
FieldName: "field-bool",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_BoolData{
BoolData: &schemapb.BoolArray{
Data: []bool{true, true, true},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData1)
fieldData2 := schemapb.FieldData{
Type: schemapb.DataType_Int8,
FieldName: "field-int8",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{
IntData: &schemapb.IntArray{
Data: []int32{0, 1, 2},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData2)
fieldData3 := schemapb.FieldData{
Type: schemapb.DataType_Int16,
FieldName: "field-int16",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{
IntData: &schemapb.IntArray{
Data: []int32{0, 1, 2},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData3)
fieldData4 := schemapb.FieldData{
Type: schemapb.DataType_Int32,
FieldName: "field-int32",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{
IntData: &schemapb.IntArray{
Data: []int32{0, 1, 2},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData4)
fieldData5 := schemapb.FieldData{
Type: schemapb.DataType_Float,
FieldName: "field-float",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_FloatData{
FloatData: &schemapb.FloatArray{
Data: []float32{0, 1, 2},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData5)
fieldData6 := schemapb.FieldData{
Type: schemapb.DataType_Double,
FieldName: "field-double",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_DoubleData{
DoubleData: &schemapb.DoubleArray{
Data: []float64{0, 1, 2},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData6)
fieldData7 := schemapb.FieldData{
Type: schemapb.DataType_String,
FieldName: "field-string",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{"0", "1", "2"},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData7)
fieldData8 := schemapb.FieldData{
Type: schemapb.DataType_VarChar,
FieldName: "field-varchar",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{"0", "1", "2"},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData8)
fieldData9 := schemapb.FieldData{
Type: schemapb.DataType_JSON,
FieldName: "field-json",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_JsonData{
JsonData: &schemapb.JSONArray{
Data: [][]byte{[]byte(`{"XXX": 0}`), []byte(`{"XXX": 0}`), []byte(`{"XXX": 0}`)},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData9)
fieldData10 := schemapb.FieldData{
Type: schemapb.DataType_Array,
FieldName: "field-array",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_ArrayData{
ArrayData: &schemapb.ArrayArray{
Data: []*schemapb.ScalarField{
{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true}}}},
{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true}}}},
{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true}}}},
},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldData11 := schemapb.FieldData{
Type: schemapb.DataType_Int64,
FieldName: "field-int64",
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{
LongData: &schemapb.LongArray{
Data: []int64{0, 1, 2},
},
},
},
},
ValidData: []bool{true, false, true},
IsDynamic: false,
}
fieldDatas = append(fieldDatas, &fieldData11)
switch firstFieldType {
case schemapb.DataType_None:
return fieldDatas
case schemapb.DataType_Bool:
return []*schemapb.FieldData{&fieldData1}
case schemapb.DataType_Int8:
return []*schemapb.FieldData{&fieldData2}
case schemapb.DataType_Int16:
return []*schemapb.FieldData{&fieldData3}
case schemapb.DataType_Int32:
return []*schemapb.FieldData{&fieldData4}
case schemapb.DataType_Float:
return []*schemapb.FieldData{&fieldData5}
case schemapb.DataType_Double:
return []*schemapb.FieldData{&fieldData6}
case schemapb.DataType_String:
return []*schemapb.FieldData{&fieldData7}
case schemapb.DataType_VarChar:
return []*schemapb.FieldData{&fieldData8}
case schemapb.DataType_BinaryVector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_FloatVector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_Float16Vector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_BFloat16Vector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_Int8Vector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_Array:
return []*schemapb.FieldData{&fieldData10}
case schemapb.DataType_JSON:
return []*schemapb.FieldData{&fieldData9}
case schemapb.DataType_SparseFloatVector:
vectorField := generateVectorFieldData(firstFieldType)
return []*schemapb.FieldData{&vectorField}
case schemapb.DataType_Int64:
return []*schemapb.FieldData{&fieldData11}
default:
return []*schemapb.FieldData{
{
FieldName: "wrong-field-type",
Type: firstFieldType,
},
}
}
}
func newSearchResult(results []map[string]interface{}) []map[string]interface{} {
for i, result := range results {
result["field-bool"] = true
result["field-int8"] = int8(i)
result["field-int16"] = int16(i)
result["field-int32"] = int32(i)
result["field-float"] = float32(i)
result["field-double"] = float64(i)
result["field-varchar"] = strconv.Itoa(i)
result["field-string"] = strconv.Itoa(i)
result["field-json"] = []byte(`{"XXX": 0}`)
result["field-geometry"] = `POINT (30.123 -10.456)`
result["field-array"] = []bool{true}
result["array-bool"] = []bool{true}
result["array-int8"] = []int32{0}
result["array-int16"] = []int32{0}
result["array-int32"] = []int32{0}
result["array-int64"] = []int64{0}
result["array-float"] = []float32{0}
result["array-double"] = []float64{0}
result["array-varchar"] = []string{""}
result["XXX"] = float64(i)
result["YYY"] = strconv.Itoa(i)
results[i] = result
}
return results
}
func newCollectionSchemaWithArray(coll *schemapb.CollectionSchema) *schemapb.CollectionSchema {
fieldSchema1 := schemapb.FieldSchema{
Name: "array-bool",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Bool,
}
coll.Fields = append(coll.Fields, &fieldSchema1)
fieldSchema2 := schemapb.FieldSchema{
Name: "array-int8",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int8,
}
coll.Fields = append(coll.Fields, &fieldSchema2)
fieldSchema3 := schemapb.FieldSchema{
Name: "array-int16",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int16,
}
coll.Fields = append(coll.Fields, &fieldSchema3)
fieldSchema4 := schemapb.FieldSchema{
Name: "array-int32",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int32,
}
coll.Fields = append(coll.Fields, &fieldSchema4)
fieldSchema5 := schemapb.FieldSchema{
Name: "array-int64",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Int64,
}
coll.Fields = append(coll.Fields, &fieldSchema5)
fieldSchema6 := schemapb.FieldSchema{
Name: "array-float",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Float,
}
coll.Fields = append(coll.Fields, &fieldSchema6)
fieldSchema7 := schemapb.FieldSchema{
Name: "array-double",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_Double,
}
coll.Fields = append(coll.Fields, &fieldSchema7)
fieldSchema8 := schemapb.FieldSchema{
Name: "array-varchar",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_VarChar,
}
coll.Fields = append(coll.Fields, &fieldSchema8)
return coll
}
func newRowsWithArray(results []map[string]interface{}) []map[string]interface{} {
for i, result := range results {
result["array-bool"] = &schemapb.ScalarField{
Data: &schemapb.ScalarField_BoolData{
BoolData: &schemapb.BoolArray{
Data: []bool{true},
},
},
}
result["array-int8"] = &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{
IntData: &schemapb.IntArray{
Data: []int32{0},
},
},
}
result["array-int16"] = &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{
IntData: &schemapb.IntArray{
Data: []int32{0},
},
},
}
result["array-int32"] = &schemapb.ScalarField{
Data: &schemapb.ScalarField_IntData{
IntData: &schemapb.IntArray{
Data: []int32{0},
},
},
}
result["array-int64"] = &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{
LongData: &schemapb.LongArray{
Data: []int64{0},
},
},
}
result["array-float"] = &schemapb.ScalarField{
Data: &schemapb.ScalarField_FloatData{
FloatData: &schemapb.FloatArray{
Data: []float32{0},
},
},
}
result["array-double"] = &schemapb.ScalarField{
Data: &schemapb.ScalarField_DoubleData{
DoubleData: &schemapb.DoubleArray{
Data: []float64{0},
},
},
}
result["array-varchar"] = &schemapb.ScalarField{
Data: &schemapb.ScalarField_StringData{
StringData: &schemapb.StringArray{
Data: []string{""},
},
},
}
results[i] = result
}
return results
}
func TestArray(t *testing.T) {
body, _ := generateRequestBody(schemapb.DataType_Int64)
collectionSchema := generateCollectionSchema(schemapb.DataType_Int64, false, true)
err, rows, validRows := checkAndSetData(body, collectionSchema, false)
assert.Equal(t, nil, err)
assert.Equal(t, 0, len(validRows))
assert.Equal(t, true, compareRows(rows, generateRawRows(schemapb.DataType_Int64), compareRow))
data, err := anyToColumns(rows, validRows, collectionSchema, true, false)
assert.Equal(t, nil, err)
assert.Equal(t, len(collectionSchema.Fields), len(data))
body, _ = generateRequestBodyWithArray(schemapb.DataType_Int64)
collectionSchema = newCollectionSchemaWithArray(generateCollectionSchema(schemapb.DataType_Int64, false, true))
err, rows, validRows = checkAndSetData(body, collectionSchema, false)
assert.Equal(t, nil, err)
assert.Equal(t, 0, len(validRows))
assert.Equal(t, true, compareRows(rows, newRowsWithArray(generateRawRows(schemapb.DataType_Int64)), compareRow))
data, err = anyToColumns(rows, validRows, collectionSchema, true, false)
assert.Equal(t, nil, err)
assert.Equal(t, len(collectionSchema.Fields), len(data))
}
func TestVector(t *testing.T) {
floatVector := "vector-float"
binaryVector := "vector-binary"
float16Vector := "vector-float16"
bfloat16Vector := "vector-bfloat16"
sparseFloatVector := "vector-sparse-float"
int8Vector := "vector-int8"
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},
int8Vector: []int8{1, 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},
int8Vector: []int8{2, 22},
},
{
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},
int8Vector: []int8{3, 33},
},
{
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},
int8Vector: []int8{4, 44},
},
{
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},
int8Vector: []int8{-128, 127},
},
}
body, err := wrapRequestBody(testcaseRows)
assert.Nil(t, err)
primaryField := generatePrimaryField(schemapb.DataType_Int64, false)
floatVectorField := generateVectorFieldSchema(schemapb.DataType_FloatVector)
floatVectorField.Name = floatVector
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
int8VectorField := generateVectorFieldSchema(schemapb.DataType_Int8Vector)
int8VectorField.Name = int8Vector
collectionSchema := &schemapb.CollectionSchema{
Name: DefaultCollectionName,
Description: "",
AutoID: false,
Fields: []*schemapb.FieldSchema{
primaryField, floatVectorField, binaryVectorField, float16VectorField, bfloat16VectorField, sparseFloatVectorField,
},
EnableDynamicField: true,
}
err, rows, validRows := checkAndSetData(body, collectionSchema, false)
assert.Equal(t, nil, err)
for i, row := range rows {
assert.Equal(t, 2, len(row[floatVector].([]float32)))
assert.Equal(t, 1, len(row[binaryVector].([]byte)))
if fv, ok := testcaseRows[i][float16Vector].([]float32); ok {
assert.Equal(t, fv, row[float16Vector].([]float32))
} else if iv, ok := testcaseRows[i][float16Vector].([]int64); ok {
assert.Equal(t, len(iv), len(row[float16Vector].([]float32)))
} 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, fv, row[float16Vector].([]float32))
} else if iv, ok := testcaseRows[i][bfloat16Vector].([]int64); ok {
assert.Equal(t, len(iv), len(row[bfloat16Vector].([]float32)))
} 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)))
}
assert.Equal(t, 0, len(validRows))
data, err := anyToColumns(rows, validRows, collectionSchema, true, false)
assert.Equal(t, nil, err)
assert.Equal(t, len(collectionSchema.Fields)+1, len(data))
assertError := func(field string, value interface{}) {
row := make(map[string]interface{})
for k, v := range testcaseRows[0] {
row[k] = v
}
row[field] = value
body, _ = wrapRequestBody([]map[string]interface{}{row})
err, _, _ = checkAndSetData(body, collectionSchema, false)
assert.Error(t, err)
}
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})
}
func TestBuildQueryResps(t *testing.T) {
outputFields := []string{"XXX", "YYY"}
outputFieldsList := [][]string{outputFields, {"$meta"}, {"$meta", FieldBookID, FieldBookIntro, "YYY"}}
for _, theOutputFields := range outputFieldsList {
rows, err := buildQueryResp(int64(0), theOutputFields, newFieldData(generateFieldData(), schemapb.DataType_None), generateIDs(schemapb.DataType_Int64, 3), DefaultScores, true, nil)
assert.Equal(t, nil, err)
exceptRows := newSearchResult(generateSearchResult(schemapb.DataType_Int64))
assert.Equal(t, true, compareRows(rows, exceptRows, compareRow))
}
dataTypes := []schemapb.DataType{
schemapb.DataType_FloatVector, schemapb.DataType_BinaryVector, schemapb.DataType_Float16Vector,
schemapb.DataType_BFloat16Vector, schemapb.DataType_SparseFloatVector, schemapb.DataType_Int8Vector,
schemapb.DataType_Bool, schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32,
schemapb.DataType_Float, schemapb.DataType_Double,
schemapb.DataType_String, schemapb.DataType_VarChar,
schemapb.DataType_JSON, schemapb.DataType_Array,
schemapb.DataType_Geometry,
}
for _, dateType := range dataTypes {
_, err := buildQueryResp(int64(0), outputFields, newFieldData([]*schemapb.FieldData{}, dateType), generateIDs(schemapb.DataType_Int64, 3), DefaultScores, true, nil)
assert.Equal(t, nil, err)
}
_, err := buildQueryResp(int64(0), outputFields, newFieldData([]*schemapb.FieldData{}, 1000), generateIDs(schemapb.DataType_Int64, 3), DefaultScores, true, nil)
assert.Equal(t, "the type(1000) of field(wrong-field-type) is not supported, use other sdk please", err.Error())
res, err := buildQueryResp(int64(0), outputFields, []*schemapb.FieldData{}, generateIDs(schemapb.DataType_Int64, 3), DefaultScores, true, nil)
assert.Equal(t, 3, len(res))
assert.Equal(t, nil, err)
res, err = buildQueryResp(int64(0), outputFields, []*schemapb.FieldData{}, generateIDs(schemapb.DataType_Int64, 3), DefaultScores, false, nil)
assert.Equal(t, 3, len(res))
assert.Equal(t, nil, err)
res, err = buildQueryResp(int64(0), outputFields, []*schemapb.FieldData{}, generateIDs(schemapb.DataType_VarChar, 3), DefaultScores, true, nil)
assert.Equal(t, 3, len(res))
assert.Equal(t, nil, err)
_, err = buildQueryResp(int64(0), outputFields, generateFieldData(), generateIDs(schemapb.DataType_Int64, 3), DefaultScores, false, nil)
assert.Equal(t, nil, err)
// len(rows) != len(scores), didn't show distance
_, err = buildQueryResp(int64(0), outputFields, newFieldData(generateFieldData(), schemapb.DataType_None), generateIDs(schemapb.DataType_Int64, 3), []float32{0.01, 0.04}, true, nil)
assert.Equal(t, nil, err)
}
func TestConvertConsistencyLevel(t *testing.T) {
consistencyLevel, useDefaultConsistency, err := convertConsistencyLevel("")
assert.Equal(t, nil, err)
assert.Equal(t, consistencyLevel, commonpb.ConsistencyLevel_Bounded)
assert.Equal(t, true, useDefaultConsistency)
consistencyLevel, useDefaultConsistency, err = convertConsistencyLevel("Strong")
assert.Equal(t, nil, err)
assert.Equal(t, consistencyLevel, commonpb.ConsistencyLevel_Strong)
assert.Equal(t, false, useDefaultConsistency)
_, _, err = convertConsistencyLevel("test")
assert.NotNil(t, err)
}
func TestConvertToExtraParams(t *testing.T) {
indexParams := IndexParam{
MetricType: "L2",
IndexType: "IVF_FLAT",
Params: map[string]interface{}{
"nlist": 128,
},
}
params, err := convertToExtraParams(indexParams)
assert.Equal(t, nil, err)
assert.Equal(t, 3, len(params))
for _, pair := range params {
if pair.Key == common.MetricTypeKey {
assert.Equal(t, "L2", pair.Value)
}
if pair.Key == common.IndexTypeKey {
assert.Equal(t, "IVF_FLAT", pair.Value)
}
if pair.Key == common.ParamsKey {
assert.Equal(t, string("{\"nlist\":128}"), pair.Value)
}
}
}
func TestGenerateExpressionTemplate(t *testing.T) {
var mixedList []interface{}
var mixedAns [][]byte
mixedList = append(mixedList, float64(1))
mixedList = append(mixedList, "10")
mixedList = append(mixedList, true)
val, _ := json.Marshal(1)
mixedAns = append(mixedAns, val)
val, _ = json.Marshal("10")
mixedAns = append(mixedAns, val)
val, _ = json.Marshal(true)
mixedAns = append(mixedAns, val)
// all passed number is float64 type, so all the number type has convert to float64
expressionTemplates := []map[string]interface{}{
{"str": "10"},
{"min": float64(1), "max": float64(10)},
{"bool": true},
{"float64": 1.1},
{"int64": float64(1)},
{"list_of_str": []interface{}{"1", "10", "100"}},
{"list_of_bool": []interface{}{true, false, true}},
{"list_of_float": []interface{}{1.1, 10.1, 100.1}},
{"list_of_int": []interface{}{float64(1), float64(10), float64(100)}},
{"list_of_json": mixedList},
}
ans := []map[string]*schemapb.TemplateValue{
{
"str": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_StringVal{
StringVal: "10",
},
},
},
{
"min": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_Int64Val{
Int64Val: 1,
},
},
"max": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_Int64Val{
Int64Val: 10,
},
},
},
{
"bool": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_BoolVal{
BoolVal: true,
},
},
},
{
"float64": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_FloatVal{
FloatVal: 1.1,
},
},
},
{
"int64": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_Int64Val{
Int64Val: 1,
},
},
},
{
"list_of_str": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_ArrayVal{
ArrayVal: &schemapb.TemplateArrayValue{
Data: &schemapb.TemplateArrayValue_StringData{
StringData: &schemapb.StringArray{
Data: []string{"1", "10", "100"},
},
},
},
},
},
},
{
"list_of_bool": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_ArrayVal{
ArrayVal: &schemapb.TemplateArrayValue{
Data: &schemapb.TemplateArrayValue_BoolData{
BoolData: &schemapb.BoolArray{
Data: []bool{true, false, true},
},
},
},
},
},
},
{
"list_of_float": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_ArrayVal{
ArrayVal: &schemapb.TemplateArrayValue{
Data: &schemapb.TemplateArrayValue_DoubleData{
DoubleData: &schemapb.DoubleArray{
Data: []float64{1.1, 10.1, 100.1},
},
},
},
},
},
},
{
"list_of_int": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_ArrayVal{
ArrayVal: &schemapb.TemplateArrayValue{
Data: &schemapb.TemplateArrayValue_LongData{
LongData: &schemapb.LongArray{
Data: []int64{1, 10, 100},
},
},
},
},
},
},
{
"list_of_json": &schemapb.TemplateValue{
Val: &schemapb.TemplateValue_ArrayVal{
ArrayVal: &schemapb.TemplateArrayValue{
Data: &schemapb.TemplateArrayValue_JsonData{
JsonData: &schemapb.JSONArray{
Data: mixedAns,
},
},
},
},
},
},
}
for i, template := range expressionTemplates {
actual := generateExpressionTemplate(template)
assert.Equal(t, actual, ans[i])
}
}
func TestGenerateSearchParams(t *testing.T) {
t.Run("searchParams.params must be a dict", func(t *testing.T) {
reqSearchParams := map[string]interface{}{"params": 0}
_, err := generateSearchParams(reqSearchParams)
assert.NotNil(t, err)
})
t.Run("ambiguous parameter", func(t *testing.T) {
reqSearchParams := map[string]interface{}{"radius": 100, "params": map[string]interface{}{"radius": 10}}
_, err := generateSearchParams(reqSearchParams)
assert.NotNil(t, err)
})
t.Run("no ambiguous parameter", func(t *testing.T) {
reqSearchParams := map[string]interface{}{"radius": 10, "params": map[string]interface{}{"radius": 10.0}}
_, err := generateSearchParams(reqSearchParams)
assert.Nil(t, err)
reqSearchParams = map[string]interface{}{"radius": 10.0, "params": map[string]interface{}{"radius": 10}}
_, err = generateSearchParams(reqSearchParams)
assert.Nil(t, err)
reqSearchParams = map[string]interface{}{"radius": 10, "params": map[string]interface{}{"radius": 10}}
searchParams, err := generateSearchParams(reqSearchParams)
assert.Equal(t, 3, len(searchParams))
assert.Nil(t, err)
for _, kvs := range searchParams {
if kvs.Key == "radius" {
assert.Equal(t, "10", kvs.Value)
}
if kvs.Key == "params" {
var paramsMap map[string]interface{}
err := json.Unmarshal([]byte(kvs.Value), &paramsMap)
assert.Nil(t, err)
assert.Equal(t, 1, len(paramsMap))
assert.Equal(t, paramsMap["radius"], float64(10))
}
}
})
t.Run("old format", func(t *testing.T) {
reqSearchParams := map[string]interface{}{"metric_type": "L2", "params": map[string]interface{}{"radius": 10}}
searchParams, err := generateSearchParams(reqSearchParams)
assert.Nil(t, err)
assert.Equal(t, 3, len(searchParams))
for _, kvs := range searchParams {
if kvs.Key == "metric_type" {
assert.Equal(t, "L2", kvs.Value)
}
if kvs.Key == "params" {
var paramsMap map[string]interface{}
err := json.Unmarshal([]byte(kvs.Value), &paramsMap)
assert.Nil(t, err)
assert.Equal(t, 2, len(paramsMap))
assert.Equal(t, paramsMap["radius"], float64(10))
assert.Equal(t, paramsMap["metric_type"], "L2")
}
}
})
t.Run("new format", func(t *testing.T) {
reqSearchParams := map[string]interface{}{"metric_type": "L2", "radius": 10}
searchParams, err := generateSearchParams(reqSearchParams)
assert.Nil(t, err)
assert.Equal(t, 4, len(searchParams))
for _, kvs := range searchParams {
if kvs.Key == "metric_type" {
assert.Equal(t, "L2", kvs.Value)
}
if kvs.Key == "radius" {
assert.Equal(t, "10", kvs.Value)
}
if kvs.Key == "params" {
var paramsMap map[string]interface{}
err := json.Unmarshal([]byte(kvs.Value), &paramsMap)
assert.Nil(t, err)
assert.Equal(t, 2, len(paramsMap))
assert.Equal(t, paramsMap["radius"], float64(10))
assert.Equal(t, paramsMap["metric_type"], "L2")
}
}
})
}
func TestGenFunctionSchem(t *testing.T) {
{
funcSchema := &FunctionSchema{
FunctionName: "test",
Description: "",
FunctionType: "unknow",
InputFieldNames: []string{"test"},
}
_, err := genFunctionSchema(context.Background(), funcSchema)
assert.ErrorContains(t, err, "Unsupported function type:")
}
{
funcSchema := &FunctionSchema{
FunctionName: "test",
Description: "",
FunctionType: "Rerank",
InputFieldNames: []string{"test"},
}
_, err := genFunctionSchema(context.Background(), funcSchema)
assert.NoError(t, err)
}
{
funcSchema := &FunctionSchema{
FunctionName: "test",
Description: "",
FunctionType: "Rerank",
InputFieldNames: []string{"test"},
Params: map[string]interface{}{
"test": []string{"test", "test2"},
"test2": map[string]interface{}{
"test3": "test4",
},
"test3": []int{1, 2, 3},
},
}
_, err := genFunctionSchema(context.Background(), funcSchema)
assert.NoError(t, err)
}
}
func TestGenFunctionScore(t *testing.T) {
{
fScore := FunctionScore{}
funcSchema := FunctionSchema{
FunctionName: "test",
Description: "",
FunctionType: "unknow",
InputFieldNames: []string{"test"},
}
fScore.Functions = append(fScore.Functions, funcSchema)
_, err := genFunctionScore(context.Background(), &fScore)
assert.ErrorContains(t, err, "Unsupported function typ")
}
{
fScore := FunctionScore{}
funcSchema := FunctionSchema{
FunctionName: "test",
Description: "",
FunctionType: "Rerank",
InputFieldNames: []string{"test"},
}
fScore.Functions = append(fScore.Functions, funcSchema)
fScore.Params = map[string]interface{}{"testStr": "test", "testInt": 6, "testBool": true}
_, err := genFunctionScore(context.Background(), &fScore)
assert.NoError(t, err)
}
}
Reference in New Issue View Git Blame Copy Permalink