wehub-resource-sync
498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
4811 lines
154 KiB
Go
4811 lines
154 KiB
Go
// Licensed to the LF AI & Data foundation under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package httpserver
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math"
|
|
"net/http/httptest"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/gjson"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
|
"github.com/milvus-io/milvus/internal/json"
|
|
"github.com/milvus-io/milvus/pkg/v3/common"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/v3/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 {
|
|
switch field.DataType {
|
|
case schemapb.DataType_VarChar:
|
|
fields = append(fields, field)
|
|
case 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
|
|
req.Data, _, err = 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
|
|
req.Data, _, err = 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
|
|
req.Data, _, err = 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("insert,autoid==true,allow_insert_auto_id=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)
|
|
coll.Properties = append(coll.Properties, &commonpb.KeyValuePair{
|
|
Key: common.AllowInsertAutoIDKey,
|
|
Value: "true",
|
|
})
|
|
var err error
|
|
req.Data, _, err = 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"])
|
|
t.Log(req.Data)
|
|
_, err = anyToColumns(req.Data, nil, coll, true, false)
|
|
assert.NoError(t, err)
|
|
})
|
|
|
|
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
|
|
req.Data, _, err = 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
|
|
req.Data, _, err = 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
|
|
req.Data, _, err = 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
|
|
req.Data, _, err = 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"])
|
|
})
|
|
|
|
t.Run("function output field not provided in any row", func(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test_collection",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "id",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "vec",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: common.DimKey, Value: "2"},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 102,
|
|
Name: "fn_out",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsFunctionOutput: true,
|
|
},
|
|
},
|
|
}
|
|
rows := []map[string]interface{}{
|
|
{"id": int64(1), "vec": []float32{0.1, 0.2}},
|
|
{"id": int64(2), "vec": []float32{0.3, 0.4}},
|
|
}
|
|
fieldsData, err := anyToColumns(rows, nil, schema, true, false)
|
|
assert.NoError(t, err)
|
|
fieldNames := make(map[string]bool)
|
|
for _, fd := range fieldsData {
|
|
fieldNames[fd.FieldName] = true
|
|
}
|
|
assert.True(t, fieldNames["id"])
|
|
assert.True(t, fieldNames["vec"])
|
|
assert.False(t, fieldNames["fn_out"])
|
|
})
|
|
|
|
t.Run("function output field provided in all rows", func(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test_collection",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "id",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "vec",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: common.DimKey, Value: "2"},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 102,
|
|
Name: "fn_out",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsFunctionOutput: true,
|
|
},
|
|
},
|
|
}
|
|
rows := []map[string]interface{}{
|
|
{"id": int64(1), "vec": []float32{0.1, 0.2}, "fn_out": int64(10)},
|
|
{"id": int64(2), "vec": []float32{0.3, 0.4}, "fn_out": int64(20)},
|
|
}
|
|
fieldsData, err := anyToColumns(rows, nil, schema, true, false)
|
|
assert.NoError(t, err)
|
|
fieldNames := make(map[string]bool)
|
|
for _, fd := range fieldsData {
|
|
fieldNames[fd.FieldName] = true
|
|
}
|
|
assert.True(t, fieldNames["fn_out"])
|
|
})
|
|
|
|
t.Run("function output field provided in row 0 but missing in later row", func(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test_collection",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "id",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "vec",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: common.DimKey, Value: "2"},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 102,
|
|
Name: "fn_out",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsFunctionOutput: true,
|
|
},
|
|
},
|
|
}
|
|
rows := []map[string]interface{}{
|
|
{"id": int64(1), "vec": []float32{0.1, 0.2}, "fn_out": int64(10)},
|
|
{"id": int64(2), "vec": []float32{0.3, 0.4}}, // fn_out missing
|
|
}
|
|
_, err := anyToColumns(rows, nil, schema, true, false)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "does not has field fn_out")
|
|
})
|
|
|
|
t.Run("function output field missing in row 0 but provided in later row", func(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test_collection",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 100,
|
|
Name: "id",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "vec",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: common.DimKey, Value: "2"},
|
|
},
|
|
},
|
|
{
|
|
FieldID: 102,
|
|
Name: "fn_out",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsFunctionOutput: true,
|
|
},
|
|
},
|
|
}
|
|
rows := []map[string]interface{}{
|
|
{"id": int64(1), "vec": []float32{0.1, 0.2}}, // fn_out missing
|
|
{"id": int64(2), "vec": []float32{0.3, 0.4}, "fn_out": int64(20)}, // fn_out provided
|
|
}
|
|
// row 0 doesn't have fn_out but row 1 does, column is allocated,
|
|
// so row 0 hits the "does not has field" error
|
|
_, err := anyToColumns(rows, nil, schema, true, false)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "does not has field fn_out")
|
|
})
|
|
}
|
|
|
|
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,
|
|
})
|
|
data, validData, err := 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,
|
|
})
|
|
data, validData, err := 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,
|
|
})
|
|
data, validData, err := 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,
|
|
})
|
|
data, validData, err := 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]}]")
|
|
data, validData, err := 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 TestInsertWithNullableVectorFields(t *testing.T) {
|
|
testcases := []struct {
|
|
name string
|
|
dataType schemapb.DataType
|
|
vectorValue interface{}
|
|
checkData func(*testing.T, *schemapb.FieldData)
|
|
}{
|
|
{
|
|
name: "float vector",
|
|
dataType: schemapb.DataType_FloatVector,
|
|
vectorValue: []float32{0.1, 0.2},
|
|
checkData: func(t *testing.T, fieldData *schemapb.FieldData) {
|
|
assert.Equal(t, []float32{0.1, 0.2}, fieldData.GetVectors().GetFloatVector().GetData())
|
|
},
|
|
},
|
|
{
|
|
name: "binary vector",
|
|
dataType: schemapb.DataType_BinaryVector,
|
|
vectorValue: []byte{1},
|
|
checkData: func(t *testing.T, fieldData *schemapb.FieldData) {
|
|
assert.Equal(t, []byte{1}, fieldData.GetVectors().GetBinaryVector())
|
|
},
|
|
},
|
|
{
|
|
name: "float16 vector",
|
|
dataType: schemapb.DataType_Float16Vector,
|
|
vectorValue: []float32{0.1, 0.2},
|
|
checkData: func(t *testing.T, fieldData *schemapb.FieldData) {
|
|
assert.Len(t, fieldData.GetVectors().GetFloat16Vector(), 4)
|
|
},
|
|
},
|
|
{
|
|
name: "bfloat16 vector",
|
|
dataType: schemapb.DataType_BFloat16Vector,
|
|
vectorValue: []float32{0.1, 0.2},
|
|
checkData: func(t *testing.T, fieldData *schemapb.FieldData) {
|
|
assert.Len(t, fieldData.GetVectors().GetBfloat16Vector(), 4)
|
|
},
|
|
},
|
|
{
|
|
name: "sparse float vector",
|
|
dataType: schemapb.DataType_SparseFloatVector,
|
|
vectorValue: map[uint32]float32{1: 0.1, 2: 0.2},
|
|
checkData: func(t *testing.T, fieldData *schemapb.FieldData) {
|
|
assert.Len(t, fieldData.GetVectors().GetSparseFloatVector().GetContents(), 1)
|
|
},
|
|
},
|
|
{
|
|
name: "int8 vector",
|
|
dataType: schemapb.DataType_Int8Vector,
|
|
vectorValue: []int8{1, 2},
|
|
checkData: func(t *testing.T, fieldData *schemapb.FieldData) {
|
|
assert.Equal(t, []byte{1, 2}, fieldData.GetVectors().GetInt8Vector())
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, testcase := range testcases {
|
|
t.Run(testcase.name, func(t *testing.T) {
|
|
primaryField := generatePrimaryField(schemapb.DataType_Int64, false)
|
|
vectorField := generateVectorFieldSchema(testcase.dataType)
|
|
vectorField.Name = FieldBookIntro
|
|
vectorField.Nullable = true
|
|
coll := &schemapb.CollectionSchema{
|
|
Name: DefaultCollectionName,
|
|
Fields: []*schemapb.FieldSchema{
|
|
primaryField,
|
|
vectorField,
|
|
},
|
|
}
|
|
body, err := wrapRequestBody([]map[string]interface{}{
|
|
{FieldBookID: int64(1), FieldBookIntro: nil},
|
|
{FieldBookID: int64(2), FieldBookIntro: testcase.vectorValue},
|
|
{FieldBookID: int64(3), FieldBookIntro: nil},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
rows, validData, err := checkAndSetData(body, coll, false)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []bool{false, true, false}, validData[FieldBookIntro])
|
|
|
|
fieldsData, err := anyToColumns(rows, validData, coll, true, false)
|
|
assert.NoError(t, err)
|
|
|
|
var vectorFieldData *schemapb.FieldData
|
|
for _, fieldData := range fieldsData {
|
|
if fieldData.GetFieldName() == FieldBookIntro {
|
|
vectorFieldData = fieldData
|
|
break
|
|
}
|
|
}
|
|
assert.NotNil(t, vectorFieldData)
|
|
assert.Equal(t, []bool{false, true, false}, vectorFieldData.GetValidData())
|
|
testcase.checkData(t, vectorFieldData)
|
|
})
|
|
}
|
|
}
|
|
|
|
func getFieldDataByName(fieldsData []*schemapb.FieldData, fieldName string) *schemapb.FieldData {
|
|
for _, fieldData := range fieldsData {
|
|
if fieldData.GetFieldName() == fieldName {
|
|
return fieldData
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func TestPartialUpdateWithNullableExplicitNull(t *testing.T) {
|
|
t.Run("nullable scalar all null is kept as update", func(t *testing.T) {
|
|
coll := &schemapb.CollectionSchema{
|
|
Name: DefaultCollectionName,
|
|
Fields: []*schemapb.FieldSchema{
|
|
generatePrimaryField(schemapb.DataType_Int64, false),
|
|
{
|
|
Name: "nullable",
|
|
FieldID: common.StartOfUserFieldID + 1,
|
|
DataType: schemapb.DataType_Int64,
|
|
Nullable: true,
|
|
},
|
|
},
|
|
}
|
|
body, err := wrapRequestBody([]map[string]interface{}{
|
|
{FieldBookID: int64(1), "nullable": nil},
|
|
{FieldBookID: int64(2), "nullable": nil},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
rows, validData, err := checkAndSetData(body, coll, true)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []bool{false, false}, validData["nullable"])
|
|
|
|
fieldsData, err := anyToColumns(rows, validData, coll, false, true)
|
|
assert.NoError(t, err)
|
|
nullableField := getFieldDataByName(fieldsData, "nullable")
|
|
assert.NotNil(t, nullableField)
|
|
assert.Equal(t, []bool{false, false}, nullableField.GetValidData())
|
|
assert.Empty(t, nullableField.GetScalars().GetLongData().GetData())
|
|
})
|
|
|
|
t.Run("nullable scalar mixed null and value is kept as compact update", func(t *testing.T) {
|
|
coll := &schemapb.CollectionSchema{
|
|
Name: DefaultCollectionName,
|
|
Fields: []*schemapb.FieldSchema{
|
|
generatePrimaryField(schemapb.DataType_Int64, false),
|
|
{
|
|
Name: "nullable",
|
|
FieldID: common.StartOfUserFieldID + 1,
|
|
DataType: schemapb.DataType_Int64,
|
|
Nullable: true,
|
|
},
|
|
},
|
|
}
|
|
body, err := wrapRequestBody([]map[string]interface{}{
|
|
{FieldBookID: int64(1), "nullable": nil},
|
|
{FieldBookID: int64(2), "nullable": int64(20)},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
rows, validData, err := checkAndSetData(body, coll, true)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []bool{false, true}, validData["nullable"])
|
|
|
|
fieldsData, err := anyToColumns(rows, validData, coll, false, true)
|
|
assert.NoError(t, err)
|
|
nullableField := getFieldDataByName(fieldsData, "nullable")
|
|
assert.NotNil(t, nullableField)
|
|
assert.Equal(t, []bool{false, true}, nullableField.GetValidData())
|
|
assert.Equal(t, []int64{20}, nullableField.GetScalars().GetLongData().GetData())
|
|
})
|
|
|
|
t.Run("missing nullable field is skipped for partial update", func(t *testing.T) {
|
|
coll := &schemapb.CollectionSchema{
|
|
Name: DefaultCollectionName,
|
|
Fields: []*schemapb.FieldSchema{
|
|
generatePrimaryField(schemapb.DataType_Int64, false),
|
|
{
|
|
Name: "nullable",
|
|
FieldID: common.StartOfUserFieldID + 1,
|
|
DataType: schemapb.DataType_Int64,
|
|
Nullable: true,
|
|
},
|
|
},
|
|
}
|
|
body, err := wrapRequestBody([]map[string]interface{}{
|
|
{FieldBookID: int64(1)},
|
|
{FieldBookID: int64(2)},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
rows, validData, err := checkAndSetData(body, coll, true)
|
|
assert.NoError(t, err)
|
|
assert.NotContains(t, validData, "nullable")
|
|
|
|
fieldsData, err := anyToColumns(rows, validData, coll, false, true)
|
|
assert.NoError(t, err)
|
|
assert.Nil(t, getFieldDataByName(fieldsData, "nullable"))
|
|
})
|
|
|
|
t.Run("mixed missing and null nullable field is rejected for partial update", func(t *testing.T) {
|
|
coll := &schemapb.CollectionSchema{
|
|
Name: DefaultCollectionName,
|
|
Fields: []*schemapb.FieldSchema{
|
|
generatePrimaryField(schemapb.DataType_Int64, false),
|
|
{
|
|
Name: "nullable",
|
|
FieldID: common.StartOfUserFieldID + 1,
|
|
DataType: schemapb.DataType_Int64,
|
|
Nullable: true,
|
|
},
|
|
},
|
|
}
|
|
body, err := wrapRequestBody([]map[string]interface{}{
|
|
{FieldBookID: int64(1)},
|
|
{FieldBookID: int64(2), "nullable": nil},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
rows, validData, err := checkAndSetData(body, coll, true)
|
|
assert.NoError(t, err)
|
|
_, err = anyToColumns(rows, validData, coll, false, true)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "column nullable has length 1, expected 2")
|
|
})
|
|
|
|
t.Run("nullable vector all null is kept as update", func(t *testing.T) {
|
|
vectorField := generateVectorFieldSchema(schemapb.DataType_FloatVector)
|
|
vectorField.Name = FieldBookIntro
|
|
vectorField.Nullable = true
|
|
coll := &schemapb.CollectionSchema{
|
|
Name: DefaultCollectionName,
|
|
Fields: []*schemapb.FieldSchema{
|
|
generatePrimaryField(schemapb.DataType_Int64, false),
|
|
vectorField,
|
|
},
|
|
}
|
|
body, err := wrapRequestBody([]map[string]interface{}{
|
|
{FieldBookID: int64(1), FieldBookIntro: nil},
|
|
{FieldBookID: int64(2), FieldBookIntro: nil},
|
|
})
|
|
assert.NoError(t, err)
|
|
|
|
rows, validData, err := checkAndSetData(body, coll, true)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []bool{false, false}, validData[FieldBookIntro])
|
|
|
|
fieldsData, err := anyToColumns(rows, validData, coll, false, true)
|
|
assert.NoError(t, err)
|
|
vectorFieldData := getFieldDataByName(fieldsData, FieldBookIntro)
|
|
assert.NotNil(t, vectorFieldData)
|
|
assert.Equal(t, []bool{false, false}, vectorFieldData.GetValidData())
|
|
assert.Empty(t, vectorFieldData.GetVectors().GetFloatVector().GetData())
|
|
assert.Equal(t, int64(2), vectorFieldData.GetVectors().GetDim())
|
|
})
|
|
}
|
|
|
|
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]}]")
|
|
data, validData, err := 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 TestBuildQueryRespWithNullableCompactFields(t *testing.T) {
|
|
t.Run("nullable vector derives logical rows from ValidData", func(t *testing.T) {
|
|
fieldData := &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.2, 0.3, 0.4},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
ValidData: []bool{true, false, true},
|
|
}
|
|
|
|
rows, err := buildQueryResp(0, []string{FieldBookIntro}, []*schemapb.FieldData{fieldData}, nil, nil, true, nil)
|
|
assert.NoError(t, err)
|
|
assert.Len(t, rows, 3)
|
|
assert.Equal(t, []float32{0.1, 0.2}, rows[0][FieldBookIntro])
|
|
assert.Nil(t, rows[1][FieldBookIntro])
|
|
assert.Equal(t, []float32{0.3, 0.4}, rows[2][FieldBookIntro])
|
|
})
|
|
|
|
t.Run("nullable vector all null keeps logical rows", func(t *testing.T) {
|
|
fieldData := &schemapb.FieldData{
|
|
Type: schemapb.DataType_FloatVector,
|
|
FieldName: FieldBookIntro,
|
|
Field: &schemapb.FieldData_Vectors{
|
|
Vectors: &schemapb.VectorField{
|
|
Dim: 2,
|
|
Data: &schemapb.VectorField_FloatVector{
|
|
FloatVector: &schemapb.FloatArray{},
|
|
},
|
|
},
|
|
},
|
|
ValidData: []bool{false, false},
|
|
}
|
|
|
|
rows, err := buildQueryResp(0, []string{FieldBookIntro}, []*schemapb.FieldData{fieldData}, nil, nil, true, nil)
|
|
assert.NoError(t, err)
|
|
assert.Len(t, rows, 2)
|
|
assert.Nil(t, rows[0][FieldBookIntro])
|
|
assert.Nil(t, rows[1][FieldBookIntro])
|
|
})
|
|
|
|
t.Run("nullable vector rejects full row payload", func(t *testing.T) {
|
|
fieldData := &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.2,
|
|
0.0, 0.0,
|
|
0.3, 0.4,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
ValidData: []bool{true, false, true},
|
|
}
|
|
|
|
_, err := buildQueryResp(0, []string{FieldBookIntro}, []*schemapb.FieldData{fieldData}, nil, nil, true, nil)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "nullable vector field")
|
|
})
|
|
|
|
t.Run("nullable dense vector rejects partial row payload", func(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
fieldData *schemapb.FieldData
|
|
}{
|
|
{
|
|
name: "float_vector",
|
|
fieldData: &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.2, 0.3}},
|
|
},
|
|
},
|
|
},
|
|
ValidData: []bool{true, false},
|
|
},
|
|
},
|
|
{
|
|
name: "binary_vector",
|
|
fieldData: &schemapb.FieldData{
|
|
Type: schemapb.DataType_BinaryVector,
|
|
FieldName: FieldBookIntro,
|
|
Field: &schemapb.FieldData_Vectors{
|
|
Vectors: &schemapb.VectorField{
|
|
Dim: 16,
|
|
Data: &schemapb.VectorField_BinaryVector{
|
|
BinaryVector: []byte{0x01, 0x02, 0x03},
|
|
},
|
|
},
|
|
},
|
|
ValidData: []bool{true, false},
|
|
},
|
|
},
|
|
{
|
|
name: "float16_vector",
|
|
fieldData: &schemapb.FieldData{
|
|
Type: schemapb.DataType_Float16Vector,
|
|
FieldName: FieldBookIntro,
|
|
Field: &schemapb.FieldData_Vectors{
|
|
Vectors: &schemapb.VectorField{
|
|
Dim: 2,
|
|
Data: &schemapb.VectorField_Float16Vector{
|
|
Float16Vector: []byte{0x01, 0x02, 0x03, 0x04, 0x05},
|
|
},
|
|
},
|
|
},
|
|
ValidData: []bool{true, false},
|
|
},
|
|
},
|
|
{
|
|
name: "bfloat16_vector",
|
|
fieldData: &schemapb.FieldData{
|
|
Type: schemapb.DataType_BFloat16Vector,
|
|
FieldName: FieldBookIntro,
|
|
Field: &schemapb.FieldData_Vectors{
|
|
Vectors: &schemapb.VectorField{
|
|
Dim: 2,
|
|
Data: &schemapb.VectorField_Bfloat16Vector{
|
|
Bfloat16Vector: []byte{0x01, 0x02, 0x03, 0x04, 0x05},
|
|
},
|
|
},
|
|
},
|
|
ValidData: []bool{true, false},
|
|
},
|
|
},
|
|
{
|
|
name: "int8_vector",
|
|
fieldData: &schemapb.FieldData{
|
|
Type: schemapb.DataType_Int8Vector,
|
|
FieldName: FieldBookIntro,
|
|
Field: &schemapb.FieldData_Vectors{
|
|
Vectors: &schemapb.VectorField{
|
|
Dim: 2,
|
|
Data: &schemapb.VectorField_Int8Vector{
|
|
Int8Vector: []byte{0x01, 0x02, 0x03},
|
|
},
|
|
},
|
|
},
|
|
ValidData: []bool{true, false},
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
_, err := buildQueryResp(0, []string{FieldBookIntro}, []*schemapb.FieldData{tc.fieldData}, nil, nil, true, nil)
|
|
require.Error(t, err)
|
|
assert.True(t,
|
|
strings.Contains(err.Error(), "row width") || strings.Contains(err.Error(), "divide the dim"),
|
|
"unexpected error: %s", err.Error())
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("nullable scalar compact data uses physical index", func(t *testing.T) {
|
|
fieldData := &schemapb.FieldData{
|
|
Type: schemapb.DataType_Int64,
|
|
FieldName: FieldWordCount,
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_LongData{
|
|
LongData: &schemapb.LongArray{Data: []int64{20}},
|
|
},
|
|
},
|
|
},
|
|
ValidData: []bool{false, true},
|
|
}
|
|
|
|
rows, err := buildQueryResp(0, []string{FieldWordCount}, []*schemapb.FieldData{fieldData}, nil, nil, true, nil)
|
|
assert.NoError(t, err)
|
|
assert.Len(t, rows, 2)
|
|
assert.Nil(t, rows[0][FieldWordCount])
|
|
assert.Equal(t, int64(20), rows[1][FieldWordCount])
|
|
})
|
|
}
|
|
|
|
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)
|
|
rows, validRows, err := 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))
|
|
rows, validRows, err = 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,
|
|
}
|
|
rows, validRows, err := 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.Contains(t, err.Error(), "the type(1000) of field(wrong-field-type) is not supported, use other sdk please")
|
|
assert.True(t, errors.Is(err, merr.ErrParameterInvalid))
|
|
|
|
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), ¶msMap)
|
|
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), ¶msMap)
|
|
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), ¶msMap)
|
|
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 TestConvertSearchAggregationReq(t *testing.T) {
|
|
req := &SearchAggregationReq{
|
|
Fields: []string{" brand "},
|
|
Size: 3,
|
|
SearchSize: 5,
|
|
Metrics: map[string]MetricAggregationReq{
|
|
" avg_price ": {Op: " avg ", FieldName: " price "},
|
|
},
|
|
Order: []AggregationOrderReq{{Key: " avg_price ", Direction: " desc "}},
|
|
TopHits: &TopHitsReq{
|
|
Size: 2,
|
|
Sort: []AggregationSortReq{{FieldName: " _score ", Direction: " asc "}},
|
|
},
|
|
SubAggregation: &SearchAggregationReq{
|
|
Fields: []string{"color"},
|
|
Size: 2,
|
|
},
|
|
}
|
|
|
|
spec, err := convertSearchAggregationReq(req)
|
|
require.NoError(t, err)
|
|
require.Equal(t, []string{"brand"}, spec.GetFields())
|
|
require.EqualValues(t, 3, spec.GetSize())
|
|
require.EqualValues(t, 5, spec.GetSearchSize())
|
|
require.Equal(t, "avg", spec.GetMetrics()["avg_price"].GetOp())
|
|
require.Equal(t, "price", spec.GetMetrics()["avg_price"].GetFieldName())
|
|
require.Equal(t, "avg_price", spec.GetOrder()[0].GetKey())
|
|
require.Equal(t, "desc", spec.GetOrder()[0].GetDirection())
|
|
require.EqualValues(t, 2, spec.GetTopHits().GetSize())
|
|
require.Equal(t, "_score", spec.GetTopHits().GetSort()[0].GetFieldName())
|
|
require.Equal(t, "color", spec.GetSubAggregation().GetFields()[0])
|
|
|
|
testCases := []struct {
|
|
name string
|
|
req *SearchAggregationReq
|
|
msg string
|
|
}{
|
|
{name: "empty fields", req: &SearchAggregationReq{Size: 1}, msg: "fields must be non-empty"},
|
|
{name: "blank field", req: &SearchAggregationReq{Fields: []string{" "}, Size: 1}, msg: "non-empty field names"},
|
|
{name: "bad size", req: &SearchAggregationReq{Fields: []string{"brand"}, Size: 0}, msg: "size must be positive"},
|
|
{name: "bad search size", req: &SearchAggregationReq{Fields: []string{"brand"}, Size: 2, SearchSize: 1}, msg: "greater than or equal to size"},
|
|
{name: "blank metric op", req: &SearchAggregationReq{Fields: []string{"brand"}, Size: 1, Metrics: map[string]MetricAggregationReq{"m": {FieldName: "price"}}}, msg: "op must be non-empty"},
|
|
{name: "blank order direction", req: &SearchAggregationReq{Fields: []string{"brand"}, Size: 1, Order: []AggregationOrderReq{{Key: "_count"}}}, msg: "direction must be non-empty"},
|
|
{name: "bad top hits size", req: &SearchAggregationReq{Fields: []string{"brand"}, Size: 1, TopHits: &TopHitsReq{}}, msg: "topHits.size must be positive"},
|
|
}
|
|
|
|
for _, testcase := range testCases {
|
|
t.Run(testcase.name, func(t *testing.T) {
|
|
_, err := convertSearchAggregationReq(testcase.req)
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), testcase.msg)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBuildSearchAggregationResp(t *testing.T) {
|
|
results := &schemapb.SearchResultData{
|
|
NumQueries: 1,
|
|
AggTopks: []int64{1},
|
|
AggBuckets: []*schemapb.AggBucket{
|
|
{
|
|
Key: []*schemapb.BucketKeyEntry{
|
|
{FieldId: 101, FieldName: "brand", Value: &schemapb.BucketKeyEntry_StringVal{StringVal: "acme"}},
|
|
{FieldId: 102, FieldName: "model_id", Value: &schemapb.BucketKeyEntry_IntVal{IntVal: 9}},
|
|
},
|
|
Count: 3,
|
|
Metrics: map[string]*schemapb.MetricValue{
|
|
"avg_price": {Value: &schemapb.MetricValue_DoubleVal{DoubleVal: 12.5}},
|
|
"stock": {Value: &schemapb.MetricValue_IntVal{IntVal: 7}},
|
|
},
|
|
Hits: []*schemapb.AggHit{
|
|
{
|
|
Pk: &schemapb.AggHit_IntPk{IntPk: 1001},
|
|
Score: 0.8,
|
|
Fields: []*schemapb.AggHitField{
|
|
{FieldId: 201, FieldName: "price", Value: &schemapb.AggHitField_IntVal{IntVal: 99}},
|
|
{FieldId: 202, FieldName: "title", Value: &schemapb.AggHitField_StringVal{StringVal: "item"}},
|
|
},
|
|
},
|
|
},
|
|
SubGroups: []*schemapb.AggBucket{
|
|
{
|
|
Key: []*schemapb.BucketKeyEntry{{FieldId: 103, FieldName: "color", Value: &schemapb.BucketKeyEntry_StringVal{StringVal: "red"}}},
|
|
Count: 1,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
resp, err := buildSearchAggregationResp(results, false, generateCollectionSchema(schemapb.DataType_Int64, false, true))
|
|
require.NoError(t, err)
|
|
require.Len(t, resp, 1)
|
|
buckets := resp[0]["buckets"].([]gin.H)
|
|
require.Len(t, buckets, 1)
|
|
bucket := buckets[0]
|
|
require.Equal(t, "3", bucket["count"])
|
|
|
|
keys := bucket["key"].([]gin.H)
|
|
require.Equal(t, "brand", keys[0]["fieldName"])
|
|
require.Equal(t, "101", keys[0]["fieldId"])
|
|
require.Equal(t, "acme", keys[0]["value"])
|
|
require.Equal(t, "9", keys[1]["value"])
|
|
|
|
metrics := bucket["metrics"].(gin.H)
|
|
require.Equal(t, 12.5, metrics["avg_price"])
|
|
require.Equal(t, "7", metrics["stock"])
|
|
|
|
hits := bucket["hits"].([]gin.H)
|
|
require.Equal(t, "1001", hits[0][FieldBookID])
|
|
require.Equal(t, float32(0.8), hits[0][HTTPReturnDistance])
|
|
require.Equal(t, "99", hits[0]["price"])
|
|
require.Equal(t, "item", hits[0]["title"])
|
|
|
|
subGroups := bucket["subGroups"].([]gin.H)
|
|
require.Equal(t, "1", subGroups[0]["count"])
|
|
|
|
_, err = buildSearchAggregationResp(&schemapb.SearchResultData{NumQueries: 1, AggBuckets: results.GetAggBuckets()}, true, nil)
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "missing agg_topks")
|
|
|
|
_, err = buildSearchAggregationResp(&schemapb.SearchResultData{AggTopks: []int64{1}, AggBuckets: results.GetAggBuckets()}, true, nil)
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "missing nq")
|
|
|
|
_, err = buildSearchAggregationResp(&schemapb.SearchResultData{NumQueries: 2, AggTopks: []int64{1}, AggBuckets: results.GetAggBuckets()}, true, nil)
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "does not match nq")
|
|
|
|
_, err = buildSearchAggregationResp(&schemapb.SearchResultData{NumQueries: 1, AggTopks: []int64{2}, AggBuckets: results.GetAggBuckets()}, true, nil)
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "does not match bucket count")
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|
|
|
|
func TestGenFunctionChains(t *testing.T) {
|
|
column := "$score"
|
|
chains, err := genFunctionChains([]FunctionChainReq{
|
|
{
|
|
Name: "l2",
|
|
Stage: "FunctionChainStageL2Rerank",
|
|
Ops: []FunctionChainOpReq{
|
|
{
|
|
Op: "map",
|
|
Outputs: []string{"new_score"},
|
|
Expr: &FunctionChainExprReq{
|
|
Name: "num_combine",
|
|
Args: []FunctionChainExprArgReq{
|
|
{Column: &column},
|
|
{Literal: 2.5},
|
|
},
|
|
Params: map[string]interface{}{
|
|
"mode": "weighted",
|
|
"weights": []interface{}{1.0, 2.0},
|
|
"nested": map[string]interface{}{"flag": true, "count": 3.0},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Op: "sort",
|
|
Inputs: []string{"new_score"},
|
|
Params: map[string]interface{}{"column": "new_score", "desc": true},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Len(t, chains, 1)
|
|
chainPB := chains[0]
|
|
assert.Equal(t, "l2", chainPB.GetName())
|
|
assert.Equal(t, schemapb.FunctionChainStage_FunctionChainStageL2Rerank, chainPB.GetStage())
|
|
require.Len(t, chainPB.GetOps(), 2)
|
|
|
|
mapOp := chainPB.GetOps()[0]
|
|
assert.Equal(t, "map", mapOp.GetOp())
|
|
assert.Equal(t, []string{"new_score"}, mapOp.GetOutputs())
|
|
require.NotNil(t, mapOp.GetExpr())
|
|
assert.Equal(t, "num_combine", mapOp.GetExpr().GetName())
|
|
assert.Equal(t, "$score", mapOp.GetExpr().GetArgs()[0].GetColumn().GetName())
|
|
assert.Equal(t, 2.5, mapOp.GetExpr().GetArgs()[1].GetLiteral().GetDoubleValue())
|
|
assert.Equal(t, "weighted", mapOp.GetExpr().GetParams()["mode"].GetStringValue())
|
|
assert.Equal(t, int64(1), mapOp.GetExpr().GetParams()["weights"].GetArrayValue().GetValues()[0].GetInt64Value())
|
|
assert.Equal(t, int64(3), mapOp.GetExpr().GetParams()["nested"].GetObjectValue().GetFields()["count"].GetInt64Value())
|
|
|
|
sortOp := chainPB.GetOps()[1]
|
|
assert.Equal(t, "sort", sortOp.GetOp())
|
|
assert.True(t, sortOp.GetParams()["desc"].GetBoolValue())
|
|
assert.Equal(t, "new_score", sortOp.GetParams()["column"].GetStringValue())
|
|
}
|
|
|
|
func TestGenFunctionChainsInvalid(t *testing.T) {
|
|
column := "field"
|
|
emptyColumn := " "
|
|
_, err := genFunctionChains([]FunctionChainReq{{Stage: "BadStage"}})
|
|
assert.ErrorContains(t, err, "unsupported function chain stage")
|
|
|
|
_, err = genFunctionChains([]FunctionChainReq{{Stage: "FunctionChainStageUnspecified"}})
|
|
assert.ErrorContains(t, err, "unsupported function chain stage")
|
|
|
|
_, err = genFunctionChains([]FunctionChainReq{{Stage: "FunctionChainStageL2Rerank", Ops: []FunctionChainOpReq{{Op: " "}}}})
|
|
assert.ErrorContains(t, err, "op name is empty")
|
|
|
|
_, err = genFunctionChains([]FunctionChainReq{{Stage: "FunctionChainStageL2Rerank", Ops: []FunctionChainOpReq{{Op: "map", Expr: &FunctionChainExprReq{Name: " "}}}}})
|
|
assert.ErrorContains(t, err, "expr name is empty")
|
|
|
|
_, err = genFunctionChains([]FunctionChainReq{{Stage: "FunctionChainStageL2Rerank", Ops: []FunctionChainOpReq{{Op: "map", Expr: &FunctionChainExprReq{Name: "expr", Args: []FunctionChainExprArgReq{{Column: &column, Literal: 1}}}}}}})
|
|
assert.ErrorContains(t, err, "exactly one of column or literal is required")
|
|
|
|
_, err = genFunctionChains([]FunctionChainReq{{Stage: "FunctionChainStageL2Rerank", Ops: []FunctionChainOpReq{{Op: "map", Expr: &FunctionChainExprReq{Name: "expr", Args: []FunctionChainExprArgReq{{}}}}}}})
|
|
assert.ErrorContains(t, err, "exactly one of column or literal is required")
|
|
|
|
_, err = genFunctionChains([]FunctionChainReq{{Stage: "FunctionChainStageL2Rerank", Ops: []FunctionChainOpReq{{Op: "map", Expr: &FunctionChainExprReq{Name: "expr", Args: []FunctionChainExprArgReq{{Column: &emptyColumn}}}}}}})
|
|
assert.ErrorContains(t, err, "column name is empty")
|
|
|
|
_, err = genFunctionChains([]FunctionChainReq{{Stage: "FunctionChainStageL2Rerank", Ops: []FunctionChainOpReq{{Op: "map", Params: map[string]interface{}{"bad": nil}}}}})
|
|
assert.ErrorContains(t, err, "function param value is nil")
|
|
|
|
_, err = genFunctionChains([]FunctionChainReq{{Stage: "FunctionChainStageL2Rerank", Ops: []FunctionChainOpReq{{Op: "map", Params: map[string]interface{}{" ": 1.0}}}}})
|
|
assert.ErrorContains(t, err, "param name is empty")
|
|
|
|
_, err = genFunctionChains([]FunctionChainReq{{Stage: "FunctionChainStageL2Rerank", Ops: []FunctionChainOpReq{{Op: "map", Params: map[string]interface{}{"object": map[string]interface{}{" ": 1.0}}}}}})
|
|
assert.ErrorContains(t, err, "object field name is empty")
|
|
}
|
|
|
|
func TestParseUsernamePassword(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
t.Run("token with credential separator", func(t *testing.T) {
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Request = httptest.NewRequest("GET", "/", nil)
|
|
c.Request.Header.Set("Authorization", "Bearer testuser:testpass")
|
|
|
|
username, password, ok := ParseUsernamePassword(c)
|
|
assert.True(t, ok)
|
|
assert.Equal(t, "testuser", username)
|
|
assert.Equal(t, "testpass", password)
|
|
})
|
|
|
|
t.Run("token without credential separator", func(t *testing.T) {
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Request = httptest.NewRequest("GET", "/", nil)
|
|
c.Request.Header.Set("Authorization", "Bearer tokenonly")
|
|
|
|
username, password, ok := ParseUsernamePassword(c)
|
|
assert.False(t, ok)
|
|
assert.Equal(t, "", username)
|
|
assert.Equal(t, "", password)
|
|
})
|
|
|
|
t.Run("empty authorization header", func(t *testing.T) {
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Request = httptest.NewRequest("GET", "/", nil)
|
|
|
|
username, password, ok := ParseUsernamePassword(c)
|
|
assert.False(t, ok)
|
|
assert.Equal(t, "", username)
|
|
assert.Equal(t, "", password)
|
|
})
|
|
}
|
|
|
|
func TestConvertIDsToSchemapbIDs(t *testing.T) {
|
|
int64PkField := &schemapb.FieldSchema{
|
|
FieldID: common.StartOfUserFieldID,
|
|
Name: "id",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
}
|
|
|
|
varcharPkField := &schemapb.FieldSchema{
|
|
FieldID: common.StartOfUserFieldID,
|
|
Name: "id",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_VarChar,
|
|
}
|
|
|
|
t.Run("empty ids array", func(t *testing.T) {
|
|
_, err := convertIDsToSchemapbIDs([]interface{}{}, int64PkField)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "ids array cannot be empty")
|
|
})
|
|
|
|
t.Run("int64 pk with float64 values (whole numbers)", func(t *testing.T) {
|
|
// JSON numbers are decoded as float64
|
|
ids := []interface{}{float64(1), float64(2), float64(3)}
|
|
result, err := convertIDsToSchemapbIDs(ids, int64PkField)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
intIds := result.GetIntId()
|
|
assert.NotNil(t, intIds)
|
|
assert.Equal(t, []int64{1, 2, 3}, intIds.Data)
|
|
})
|
|
|
|
t.Run("int64 pk with float64 values having fractional part", func(t *testing.T) {
|
|
ids := []interface{}{float64(1.5)}
|
|
_, err := convertIDsToSchemapbIDs(ids, int64PkField)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "has fractional part")
|
|
})
|
|
|
|
t.Run("int64 pk with float64 values - second element has fractional part", func(t *testing.T) {
|
|
ids := []interface{}{float64(1), float64(2.9)}
|
|
_, err := convertIDsToSchemapbIDs(ids, int64PkField)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "index 1")
|
|
assert.Contains(t, err.Error(), "has fractional part")
|
|
})
|
|
|
|
t.Run("int64 pk with int64 values", func(t *testing.T) {
|
|
ids := []interface{}{int64(100), int64(200)}
|
|
result, err := convertIDsToSchemapbIDs(ids, int64PkField)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
intIds := result.GetIntId()
|
|
assert.NotNil(t, intIds)
|
|
assert.Equal(t, []int64{100, 200}, intIds.Data)
|
|
})
|
|
|
|
t.Run("int64 pk with int values", func(t *testing.T) {
|
|
ids := []interface{}{int(10), int(20)}
|
|
result, err := convertIDsToSchemapbIDs(ids, int64PkField)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
intIds := result.GetIntId()
|
|
assert.NotNil(t, intIds)
|
|
assert.Equal(t, []int64{10, 20}, intIds.Data)
|
|
})
|
|
|
|
t.Run("int64 pk with valid string values", func(t *testing.T) {
|
|
ids := []interface{}{"123", "456"}
|
|
result, err := convertIDsToSchemapbIDs(ids, int64PkField)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
intIds := result.GetIntId()
|
|
assert.NotNil(t, intIds)
|
|
assert.Equal(t, []int64{123, 456}, intIds.Data)
|
|
})
|
|
|
|
t.Run("int64 pk with invalid string values", func(t *testing.T) {
|
|
ids := []interface{}{"not_a_number"}
|
|
_, err := convertIDsToSchemapbIDs(ids, int64PkField)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "invalid int64 id")
|
|
})
|
|
|
|
t.Run("int64 pk with invalid type", func(t *testing.T) {
|
|
ids := []interface{}{true}
|
|
_, err := convertIDsToSchemapbIDs(ids, int64PkField)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "invalid id type")
|
|
})
|
|
|
|
t.Run("varchar pk with string values", func(t *testing.T) {
|
|
ids := []interface{}{"abc", "def", "ghi"}
|
|
result, err := convertIDsToSchemapbIDs(ids, varcharPkField)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
strIds := result.GetStrId()
|
|
assert.NotNil(t, strIds)
|
|
assert.Equal(t, []string{"abc", "def", "ghi"}, strIds.Data)
|
|
})
|
|
|
|
t.Run("varchar pk with empty string", func(t *testing.T) {
|
|
ids := []interface{}{""}
|
|
_, err := convertIDsToSchemapbIDs(ids, varcharPkField)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "empty string id")
|
|
})
|
|
|
|
t.Run("varchar pk with number values", func(t *testing.T) {
|
|
ids := []interface{}{float64(123), int64(456), int(789)}
|
|
result, err := convertIDsToSchemapbIDs(ids, varcharPkField)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
strIds := result.GetStrId()
|
|
assert.NotNil(t, strIds)
|
|
assert.Equal(t, []string{"123", "456", "789"}, strIds.Data)
|
|
})
|
|
|
|
t.Run("varchar pk with invalid type", func(t *testing.T) {
|
|
ids := []interface{}{[]int{1, 2, 3}}
|
|
_, err := convertIDsToSchemapbIDs(ids, varcharPkField)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "invalid id type")
|
|
})
|
|
|
|
t.Run("unsupported pk type", func(t *testing.T) {
|
|
boolPkField := &schemapb.FieldSchema{
|
|
FieldID: common.StartOfUserFieldID,
|
|
Name: "id",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Bool,
|
|
}
|
|
ids := []interface{}{float64(1)}
|
|
_, err := convertIDsToSchemapbIDs(ids, boolPkField)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "unsupported primary key type")
|
|
})
|
|
}
|
|
|
|
func buildStructArrayTestSchema() *schemapb.CollectionSchema {
|
|
pk := &schemapb.FieldSchema{
|
|
FieldID: 100,
|
|
Name: "id",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
}
|
|
vec := &schemapb.FieldSchema{
|
|
FieldID: 101,
|
|
Name: "vec",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: common.DimKey, Value: "4"},
|
|
},
|
|
}
|
|
subInt := &schemapb.FieldSchema{
|
|
FieldID: 103,
|
|
Name: "sub_int",
|
|
DataType: schemapb.DataType_Array,
|
|
ElementType: schemapb.DataType_Int32,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: common.MaxCapacityKey, Value: "10"},
|
|
},
|
|
}
|
|
subVec := &schemapb.FieldSchema{
|
|
FieldID: 104,
|
|
Name: "sub_vec",
|
|
DataType: schemapb.DataType_ArrayOfVector,
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: common.DimKey, Value: "4"},
|
|
{Key: common.MaxCapacityKey, Value: "10"},
|
|
},
|
|
}
|
|
structField := &schemapb.StructArrayFieldSchema{
|
|
FieldID: 102,
|
|
Name: "my_struct",
|
|
Fields: []*schemapb.FieldSchema{subInt, subVec},
|
|
}
|
|
return &schemapb.CollectionSchema{
|
|
Name: "c_test",
|
|
Fields: []*schemapb.FieldSchema{pk, vec},
|
|
StructArrayFields: []*schemapb.StructArrayFieldSchema{structField},
|
|
}
|
|
}
|
|
|
|
func TestStructArrayFieldSchemaGetProto(t *testing.T) {
|
|
ctx := context.Background()
|
|
good := StructArrayFieldSchema{
|
|
FieldName: "my_struct",
|
|
Description: "struct field",
|
|
Fields: []FieldSchema{
|
|
{
|
|
FieldName: "sub_int",
|
|
DataType: "Array",
|
|
ElementDataType: "Int32",
|
|
ElementTypeParams: map[string]interface{}{
|
|
"max_capacity": 10,
|
|
},
|
|
},
|
|
{
|
|
FieldName: "sub_vec",
|
|
DataType: "ArrayOfVector",
|
|
ElementDataType: "FloatVector",
|
|
ElementTypeParams: map[string]interface{}{
|
|
"dim": 4,
|
|
"max_capacity": 10,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
proto, err := good.GetProto(ctx)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "my_struct", proto.GetName())
|
|
assert.Len(t, proto.GetFields(), 2)
|
|
assert.Equal(t, schemapb.DataType_Array, proto.GetFields()[0].GetDataType())
|
|
assert.Equal(t, schemapb.DataType_Int32, proto.GetFields()[0].GetElementType())
|
|
assert.Equal(t, schemapb.DataType_ArrayOfVector, proto.GetFields()[1].GetDataType())
|
|
assert.Equal(t, schemapb.DataType_FloatVector, proto.GetFields()[1].GetElementType())
|
|
|
|
_, err = (&StructArrayFieldSchema{
|
|
FieldName: "bad",
|
|
Fields: []FieldSchema{
|
|
{FieldName: "raw_int", DataType: "Int32"},
|
|
},
|
|
}).GetProto(ctx)
|
|
assert.Error(t, err)
|
|
|
|
_, err = (&StructArrayFieldSchema{
|
|
FieldName: "bad_pk",
|
|
Fields: []FieldSchema{
|
|
{
|
|
FieldName: "sub_pk",
|
|
DataType: "Array",
|
|
ElementDataType: "Int32",
|
|
IsPrimary: true,
|
|
},
|
|
},
|
|
}).GetProto(ctx)
|
|
assert.Error(t, err)
|
|
|
|
_, err = (&StructArrayFieldSchema{FieldName: "empty"}).GetProto(ctx)
|
|
assert.Error(t, err)
|
|
|
|
_, err = (&StructArrayFieldSchema{
|
|
FieldName: "dup",
|
|
Fields: []FieldSchema{
|
|
{FieldName: "s", DataType: "Array", ElementDataType: "Int32"},
|
|
{FieldName: "s", DataType: "Array", ElementDataType: "Int32"},
|
|
},
|
|
}).GetProto(ctx)
|
|
assert.Error(t, err)
|
|
|
|
_, err = (&StructArrayFieldSchema{
|
|
FieldName: "bad_nullable",
|
|
Fields: []FieldSchema{
|
|
{
|
|
FieldName: "sub_null",
|
|
DataType: "Array",
|
|
ElementDataType: "Int32",
|
|
Nullable: true,
|
|
},
|
|
},
|
|
}).GetProto(ctx)
|
|
assert.Error(t, err)
|
|
|
|
_, err = (&StructArrayFieldSchema{
|
|
FieldName: "bad_default",
|
|
Fields: []FieldSchema{
|
|
{
|
|
FieldName: "sub_default",
|
|
DataType: "Array",
|
|
ElementDataType: "Int32",
|
|
DefaultValue: float64(1),
|
|
},
|
|
},
|
|
}).GetProto(ctx)
|
|
assert.Error(t, err)
|
|
|
|
_, err = (&StructArrayFieldSchema{
|
|
FieldName: "bad_part",
|
|
Fields: []FieldSchema{
|
|
{
|
|
FieldName: "sub_part",
|
|
DataType: "Array",
|
|
ElementDataType: "Int32",
|
|
IsPartitionKey: true,
|
|
},
|
|
},
|
|
}).GetProto(ctx)
|
|
assert.Error(t, err)
|
|
|
|
_, err = (&StructArrayFieldSchema{
|
|
FieldName: "bad_cluster",
|
|
Fields: []FieldSchema{
|
|
{
|
|
FieldName: "sub_cluster",
|
|
DataType: "Array",
|
|
ElementDataType: "Int32",
|
|
IsClusteringKey: true,
|
|
},
|
|
},
|
|
}).GetProto(ctx)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestParseStructArrayRowScalar(t *testing.T) {
|
|
schema := buildStructArrayTestSchema().GetStructArrayFields()[0]
|
|
raw := `[{"sub_int": 1, "sub_vec": [0.1, 0.2, 0.3, 0.4]},
|
|
{"sub_int": 2, "sub_vec": [0.5, 0.6, 0.7, 0.8]}]`
|
|
row, err := parseStructArrayRow(raw, schema)
|
|
require.NoError(t, err)
|
|
require.Len(t, row, 2)
|
|
|
|
scalar, ok := row["sub_int"].(*schemapb.ScalarField)
|
|
require.True(t, ok)
|
|
assert.Equal(t, []int32{1, 2}, scalar.GetIntData().GetData())
|
|
|
|
vecField, ok := row["sub_vec"].(*schemapb.VectorField)
|
|
require.True(t, ok)
|
|
assert.Equal(t,
|
|
[]float32{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8},
|
|
vecField.GetFloatVector().GetData())
|
|
}
|
|
|
|
func TestParseStructArrayRowMissingField(t *testing.T) {
|
|
schema := buildStructArrayTestSchema().GetStructArrayFields()[0]
|
|
_, err := parseStructArrayRow(`[{"sub_int": 1}]`, schema)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestParseStructArrayRowNotArray(t *testing.T) {
|
|
schema := buildStructArrayTestSchema().GetStructArrayFields()[0]
|
|
_, err := parseStructArrayRow(`{"sub_int": 1}`, schema)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestBuildStructArrayFieldDataRoundTrip(t *testing.T) {
|
|
schema := buildStructArrayTestSchema().GetStructArrayFields()[0]
|
|
r1, err := parseStructArrayRow(`[{"sub_int": 1, "sub_vec": [0.1, 0.2, 0.3, 0.4]},
|
|
{"sub_int": 2, "sub_vec": [0.5, 0.6, 0.7, 0.8]}]`, schema)
|
|
require.NoError(t, err)
|
|
r2, err := parseStructArrayRow(`[{"sub_int": 3, "sub_vec": [0.9, 1.0, 1.1, 1.2]}]`, schema)
|
|
require.NoError(t, err)
|
|
|
|
fd, err := buildStructArrayFieldData(schema, []structArrayRow{r1, r2})
|
|
require.NoError(t, err)
|
|
require.Equal(t, schemapb.DataType_ArrayOfStruct, fd.GetType())
|
|
subs := fd.GetStructArrays().GetFields()
|
|
require.Len(t, subs, 2)
|
|
|
|
assert.Equal(t, schemapb.DataType_Array, subs[0].GetType())
|
|
assert.Len(t, subs[0].GetScalars().GetArrayData().GetData(), 2)
|
|
assert.Equal(t, schemapb.DataType_ArrayOfVector, subs[1].GetType())
|
|
assert.Len(t, subs[1].GetVectors().GetVectorArray().GetData(), 2)
|
|
|
|
extracted0, err := extractStructArrayRow(fd, 0, buildStructArrayTestSchema())
|
|
require.NoError(t, err)
|
|
require.Len(t, extracted0, 2)
|
|
assert.EqualValues(t, int32(1), extracted0[0]["sub_int"])
|
|
assert.EqualValues(t, []float32{0.1, 0.2, 0.3, 0.4}, extracted0[0]["sub_vec"])
|
|
|
|
extracted1, err := extractStructArrayRow(fd, 1, buildStructArrayTestSchema())
|
|
require.NoError(t, err)
|
|
require.Len(t, extracted1, 1)
|
|
assert.EqualValues(t, int32(3), extracted1[0]["sub_int"])
|
|
}
|
|
|
|
func TestAnyToColumnsStructArray(t *testing.T) {
|
|
schema := buildStructArrayTestSchema()
|
|
body := []byte(`{
|
|
"data": [
|
|
{
|
|
"id": 1,
|
|
"vec": [0.1, 0.2, 0.3, 0.4],
|
|
"my_struct": [
|
|
{"sub_int": 10, "sub_vec": [1.1, 1.2, 1.3, 1.4]},
|
|
{"sub_int": 20, "sub_vec": [2.1, 2.2, 2.3, 2.4]}
|
|
]
|
|
},
|
|
{
|
|
"id": 2,
|
|
"vec": [0.5, 0.6, 0.7, 0.8],
|
|
"my_struct": [
|
|
{"sub_int": 30, "sub_vec": [3.1, 3.2, 3.3, 3.4]}
|
|
]
|
|
}
|
|
]
|
|
}`)
|
|
rows, _, err := checkAndSetData(body, schema, false)
|
|
require.NoError(t, err)
|
|
require.Len(t, rows, 2)
|
|
|
|
fds, err := anyToColumns(rows, nil, schema, true, false)
|
|
require.NoError(t, err)
|
|
|
|
var structFD *schemapb.FieldData
|
|
for _, fd := range fds {
|
|
if fd.GetType() == schemapb.DataType_ArrayOfStruct {
|
|
structFD = fd
|
|
break
|
|
}
|
|
}
|
|
require.NotNil(t, structFD)
|
|
assert.Equal(t, "my_struct", structFD.GetFieldName())
|
|
subs := structFD.GetStructArrays().GetFields()
|
|
require.Len(t, subs, 2)
|
|
|
|
arrayData := subs[0].GetScalars().GetArrayData().GetData()
|
|
require.Len(t, arrayData, 2)
|
|
assert.Equal(t, []int32{10, 20}, arrayData[0].GetIntData().GetData())
|
|
assert.Equal(t, []int32{30}, arrayData[1].GetIntData().GetData())
|
|
|
|
vecData := subs[1].GetVectors().GetVectorArray().GetData()
|
|
require.Len(t, vecData, 2)
|
|
assert.Len(t, vecData[0].GetFloatVector().GetData(), 8)
|
|
assert.Len(t, vecData[1].GetFloatVector().GetData(), 4)
|
|
}
|
|
|
|
func TestBuildQueryRespStructArrayRoundTrip(t *testing.T) {
|
|
schema := buildStructArrayTestSchema()
|
|
body := []byte(`{
|
|
"data": [
|
|
{
|
|
"id": 1,
|
|
"vec": [0.1, 0.2, 0.3, 0.4],
|
|
"my_struct": [
|
|
{"sub_int": 10, "sub_vec": [1.1, 1.2, 1.3, 1.4]},
|
|
{"sub_int": 20, "sub_vec": [2.1, 2.2, 2.3, 2.4]}
|
|
]
|
|
}
|
|
]
|
|
}`)
|
|
rows, _, err := checkAndSetData(body, schema, false)
|
|
require.NoError(t, err)
|
|
fds, err := anyToColumns(rows, nil, schema, true, false)
|
|
require.NoError(t, err)
|
|
|
|
var structFD *schemapb.FieldData
|
|
for _, fd := range fds {
|
|
if fd.GetType() == schemapb.DataType_ArrayOfStruct {
|
|
structFD = fd
|
|
break
|
|
}
|
|
}
|
|
require.NotNil(t, structFD)
|
|
|
|
resp, err := buildQueryResp(0, []string{"my_struct"}, []*schemapb.FieldData{structFD}, nil, nil, true, schema)
|
|
require.NoError(t, err)
|
|
require.Len(t, resp, 1)
|
|
|
|
blob, err := json.Marshal(resp[0]["my_struct"])
|
|
require.NoError(t, err)
|
|
var decoded []map[string]interface{}
|
|
require.NoError(t, json.Unmarshal(blob, &decoded))
|
|
require.Len(t, decoded, 2)
|
|
assert.EqualValues(t, 10, decoded[0]["sub_int"])
|
|
}
|
|
|
|
func TestIsEmbeddingListData(t *testing.T) {
|
|
assert.False(t, isEmbeddingListData(`{"data": [[0.1, 0.2, 0.3, 0.4]]}`))
|
|
assert.False(t, isEmbeddingListData(`{"data": [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]]}`))
|
|
assert.False(t, isEmbeddingListData(`{"data": ["YmFzZTY0"]}`))
|
|
assert.False(t, isEmbeddingListData(`{"data": ["YmFzZTY0", "YmFzZTY1"]}`))
|
|
|
|
assert.True(t, isEmbeddingListData(`{"data": [[[0.1, 0.2, 0.3, 0.4]]]}`))
|
|
assert.True(t, isEmbeddingListData(`{"data": [[[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]], [[0.9, 1.0, 1.1, 1.2]]]}`))
|
|
assert.True(t, isEmbeddingListData(`{"data": [["YmFzZTY0", "YmFzZTY1"]]}`))
|
|
|
|
assert.False(t, isEmbeddingListData(`{"data": []}`))
|
|
assert.False(t, isEmbeddingListData(`{"data": "not-array"}`))
|
|
assert.False(t, isEmbeddingListData(`{"data": [[]]}`))
|
|
}
|
|
|
|
func TestPrintStructArrayFieldsV2(t *testing.T) {
|
|
schema := buildStructArrayTestSchema()
|
|
schema.GetStructArrayFields()[0].Nullable = true
|
|
printed := printStructArrayFieldsV2(schema.GetStructArrayFields())
|
|
require.Len(t, printed, 1)
|
|
entry := printed[0]
|
|
assert.Equal(t, "my_struct", entry[HTTPReturnFieldName])
|
|
assert.Equal(t, schemapb.DataType_ArrayOfStruct.String(), entry[HTTPReturnFieldType])
|
|
assert.Equal(t, true, entry[HTTPReturnFieldNullable])
|
|
subs, ok := entry["fields"].([]gin.H)
|
|
require.True(t, ok)
|
|
require.Len(t, subs, 2)
|
|
assert.Equal(t, "sub_int", subs[0][HTTPReturnFieldName])
|
|
assert.Equal(t, schemapb.DataType_Array.String(), subs[0][HTTPReturnFieldType])
|
|
assert.Equal(t, schemapb.DataType_Int32.String(), subs[0][HTTPReturnFieldElementType])
|
|
assert.Equal(t, "sub_vec", subs[1][HTTPReturnFieldName])
|
|
assert.Equal(t, schemapb.DataType_FloatVector.String(), subs[1][HTTPReturnFieldElementType])
|
|
}
|
|
|
|
func TestStructArrayScalarSubFieldTypes(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
elementType schemapb.DataType
|
|
raw string
|
|
assertFn func(*testing.T, *schemapb.ScalarField)
|
|
}{
|
|
{
|
|
name: "bool",
|
|
elementType: schemapb.DataType_Bool,
|
|
raw: `[true,false]`,
|
|
assertFn: func(t *testing.T, sf *schemapb.ScalarField) {
|
|
assert.Equal(t, []bool{true, false}, sf.GetBoolData().GetData())
|
|
},
|
|
},
|
|
{
|
|
name: "int64",
|
|
elementType: schemapb.DataType_Int64,
|
|
raw: `[10,20]`,
|
|
assertFn: func(t *testing.T, sf *schemapb.ScalarField) {
|
|
assert.Equal(t, []int64{10, 20}, sf.GetLongData().GetData())
|
|
},
|
|
},
|
|
{
|
|
name: "float",
|
|
elementType: schemapb.DataType_Float,
|
|
raw: `[1.5,2.5]`,
|
|
assertFn: func(t *testing.T, sf *schemapb.ScalarField) {
|
|
assert.Equal(t, []float32{1.5, 2.5}, sf.GetFloatData().GetData())
|
|
},
|
|
},
|
|
{
|
|
name: "double",
|
|
elementType: schemapb.DataType_Double,
|
|
raw: `[1.25,2.25]`,
|
|
assertFn: func(t *testing.T, sf *schemapb.ScalarField) {
|
|
assert.Equal(t, []float64{1.25, 2.25}, sf.GetDoubleData().GetData())
|
|
},
|
|
},
|
|
{
|
|
name: "string",
|
|
elementType: schemapb.DataType_VarChar,
|
|
raw: `["red","blue"]`,
|
|
assertFn: func(t *testing.T, sf *schemapb.ScalarField) {
|
|
assert.Equal(t, []string{"red", "blue"}, sf.GetStringData().GetData())
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
sub := &schemapb.FieldSchema{Name: test.name, ElementType: test.elementType}
|
|
got, err := buildStructSubArrayScalar(sub, gjson.Parse(test.raw).Array())
|
|
require.NoError(t, err)
|
|
test.assertFn(t, got)
|
|
})
|
|
}
|
|
|
|
_, err := buildStructSubArrayScalar(&schemapb.FieldSchema{Name: "bad", ElementType: schemapb.DataType_JSON}, gjson.Parse(`[{}]`).Array())
|
|
assert.Error(t, err)
|
|
|
|
_, err = buildStructSubArrayScalar(&schemapb.FieldSchema{Name: "bad_bool", ElementType: schemapb.DataType_Bool}, gjson.Parse(`[1]`).Array())
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestStructArrayVectorSubFieldTypes(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
elementType schemapb.DataType
|
|
dim string
|
|
raw string
|
|
assertFn func(*testing.T, *schemapb.VectorField)
|
|
}{
|
|
{
|
|
name: "float16",
|
|
elementType: schemapb.DataType_Float16Vector,
|
|
dim: "2",
|
|
raw: `[[0.1,0.2],"AQIDBA=="]`,
|
|
assertFn: func(t *testing.T, vf *schemapb.VectorField) {
|
|
assert.Len(t, vf.GetFloat16Vector(), 8)
|
|
},
|
|
},
|
|
{
|
|
name: "bfloat16",
|
|
elementType: schemapb.DataType_BFloat16Vector,
|
|
dim: "2",
|
|
raw: `[[0.1,0.2],[0.3,0.4]]`,
|
|
assertFn: func(t *testing.T, vf *schemapb.VectorField) {
|
|
assert.Len(t, vf.GetBfloat16Vector(), 8)
|
|
},
|
|
},
|
|
{
|
|
name: "binary",
|
|
elementType: schemapb.DataType_BinaryVector,
|
|
dim: "16",
|
|
raw: `["AQI=","AwQ="]`,
|
|
assertFn: func(t *testing.T, vf *schemapb.VectorField) {
|
|
assert.Equal(t, []byte{1, 2, 3, 4}, vf.GetBinaryVector())
|
|
},
|
|
},
|
|
{
|
|
name: "int8",
|
|
elementType: schemapb.DataType_Int8Vector,
|
|
dim: "2",
|
|
raw: `[[1,-2],[3,4]]`,
|
|
assertFn: func(t *testing.T, vf *schemapb.VectorField) {
|
|
assert.Equal(t, []byte{1, 254, 3, 4}, vf.GetInt8Vector())
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
sub := &schemapb.FieldSchema{
|
|
Name: test.name,
|
|
ElementType: test.elementType,
|
|
TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: test.dim}},
|
|
}
|
|
got, err := buildStructSubVectorField(sub, gjson.Parse(test.raw).Array())
|
|
require.NoError(t, err)
|
|
test.assertFn(t, got)
|
|
})
|
|
}
|
|
|
|
_, err := buildStructSubVectorField(&schemapb.FieldSchema{
|
|
Name: "bad_dim",
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "2"}},
|
|
}, gjson.Parse(`[[0.1]]`).Array())
|
|
assert.Error(t, err)
|
|
|
|
_, err = buildStructSubVectorField(&schemapb.FieldSchema{
|
|
Name: "bad_type",
|
|
ElementType: schemapb.DataType_JSON,
|
|
TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "2"}},
|
|
}, gjson.Parse(`[{}]`).Array())
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestEmbeddingListPlaceholderTypes(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
elementType schemapb.DataType
|
|
body string
|
|
dim int64
|
|
placeholderType commonpb.PlaceholderType
|
|
valueLen int
|
|
}{
|
|
{
|
|
name: "float",
|
|
elementType: schemapb.DataType_FloatVector,
|
|
body: `{"data": [[[0.1,0.2],[0.3,0.4]]]}`,
|
|
dim: 2,
|
|
placeholderType: commonpb.PlaceholderType_EmbListFloatVector,
|
|
valueLen: 16,
|
|
},
|
|
{
|
|
name: "float16",
|
|
elementType: schemapb.DataType_Float16Vector,
|
|
body: `{"data": [[[0.1,0.2],[0.3,0.4]]]}`,
|
|
dim: 2,
|
|
placeholderType: commonpb.PlaceholderType_EmbListFloat16Vector,
|
|
valueLen: 8,
|
|
},
|
|
{
|
|
name: "bfloat16",
|
|
elementType: schemapb.DataType_BFloat16Vector,
|
|
body: `{"data": [[[0.1,0.2],[0.3,0.4]]]}`,
|
|
dim: 2,
|
|
placeholderType: commonpb.PlaceholderType_EmbListBFloat16Vector,
|
|
valueLen: 8,
|
|
},
|
|
{
|
|
name: "binary",
|
|
elementType: schemapb.DataType_BinaryVector,
|
|
body: `{"data": [["AQI=","AwQ="]]}`,
|
|
dim: 16,
|
|
placeholderType: commonpb.PlaceholderType_EmbListBinaryVector,
|
|
valueLen: 4,
|
|
},
|
|
{
|
|
name: "int8",
|
|
elementType: schemapb.DataType_Int8Vector,
|
|
body: `{"data": [[[1,2],[3,4]]]}`,
|
|
dim: 2,
|
|
placeholderType: commonpb.PlaceholderType_EmbListInt8Vector,
|
|
valueLen: 4,
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
got, err := convertEmbListQueries2Placeholder(test.body, test.elementType, test.dim)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, test.placeholderType, got.GetType())
|
|
require.Len(t, got.GetValues(), 1)
|
|
assert.Len(t, got.GetValues()[0], test.valueLen)
|
|
})
|
|
}
|
|
|
|
for _, test := range []struct {
|
|
name string
|
|
body string
|
|
elementType schemapb.DataType
|
|
dim int64
|
|
}{
|
|
{name: "data_not_array", body: `{"data":"bad"}`, elementType: schemapb.DataType_FloatVector, dim: 2},
|
|
{name: "empty_data", body: `{"data":[]}`, elementType: schemapb.DataType_FloatVector, dim: 2},
|
|
{name: "query_not_array", body: `{"data":[1]}`, elementType: schemapb.DataType_FloatVector, dim: 2},
|
|
{name: "empty_embedding_list", body: `{"data":[[]]}`, elementType: schemapb.DataType_FloatVector, dim: 2},
|
|
{name: "unsupported_element", body: `{"data":[[true]]}`, elementType: schemapb.DataType_Bool, dim: 2},
|
|
{name: "float_dim_mismatch", body: `{"data":[[[0.1]]]}`, elementType: schemapb.DataType_FloatVector, dim: 2},
|
|
{name: "binary_not_string", body: `{"data":[[[1,2]]]}`, elementType: schemapb.DataType_BinaryVector, dim: 16},
|
|
{name: "int8_not_array", body: `{"data":[["AQI="]]}`, elementType: schemapb.DataType_Int8Vector, dim: 2},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
_, err := convertEmbListQueries2Placeholder(test.body, test.elementType, test.dim)
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGeneratePlaceholderGroupStructArrayField(t *testing.T) {
|
|
schema := buildStructArrayTestSchema()
|
|
|
|
standard, err := generatePlaceholderGroup(context.Background(), `{"data": [[0.1,0.2,0.3,0.4]]}`, schema, "sub_vec")
|
|
require.NoError(t, err)
|
|
group := &commonpb.PlaceholderGroup{}
|
|
require.NoError(t, proto.Unmarshal(standard, group))
|
|
require.Len(t, group.GetPlaceholders(), 1)
|
|
assert.Equal(t, commonpb.PlaceholderType_FloatVector, group.GetPlaceholders()[0].GetType())
|
|
|
|
embList, err := generatePlaceholderGroup(context.Background(), `{"data": [[[0.1,0.2,0.3,0.4]]]}`, schema, "sub_vec")
|
|
require.NoError(t, err)
|
|
group.Reset()
|
|
require.NoError(t, proto.Unmarshal(embList, group))
|
|
require.Len(t, group.GetPlaceholders(), 1)
|
|
assert.Equal(t, commonpb.PlaceholderType_EmbListFloatVector, group.GetPlaceholders()[0].GetType())
|
|
}
|
|
|
|
func TestStructArrayFieldDataErrorPaths(t *testing.T) {
|
|
schema := buildStructArrayTestSchema().GetStructArrayFields()[0]
|
|
|
|
_, err := buildStructArrayFieldData(schema, nil)
|
|
assert.Error(t, err)
|
|
|
|
_, err = buildStructArrayFieldData(schema, []structArrayRow{{"sub_int": &schemapb.ScalarField{}}})
|
|
assert.Error(t, err)
|
|
|
|
_, err = buildStructArrayFieldData(schema, []structArrayRow{
|
|
{
|
|
"sub_int": int64(1),
|
|
"sub_vec": &schemapb.VectorField{},
|
|
},
|
|
})
|
|
assert.Error(t, err)
|
|
|
|
badTypeSchema := &schemapb.StructArrayFieldSchema{
|
|
Name: "bad",
|
|
Fields: []*schemapb.FieldSchema{{Name: "sub", DataType: schemapb.DataType_Bool}},
|
|
}
|
|
_, err = buildStructArrayFieldData(badTypeSchema, []structArrayRow{{"sub": true}})
|
|
assert.Error(t, err)
|
|
|
|
noDimSchema := &schemapb.StructArrayFieldSchema{
|
|
Name: "bad_dim",
|
|
Fields: []*schemapb.FieldSchema{{
|
|
Name: "sub_vec",
|
|
DataType: schemapb.DataType_ArrayOfVector,
|
|
}},
|
|
}
|
|
_, err = buildStructArrayFieldData(noDimSchema, []structArrayRow{{"sub_vec": &schemapb.VectorField{}}})
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestExtractStructArrayRowErrorPaths(t *testing.T) {
|
|
empty, err := extractStructArrayRow(&schemapb.FieldData{
|
|
Type: schemapb.DataType_ArrayOfStruct,
|
|
Field: &schemapb.FieldData_StructArrays{
|
|
StructArrays: &schemapb.StructArrayField{},
|
|
},
|
|
}, 0, buildStructArrayTestSchema())
|
|
require.NoError(t, err)
|
|
assert.Empty(t, empty)
|
|
|
|
schema := buildStructArrayTestSchema()
|
|
row, err := parseStructArrayRow(`[{"sub_int": 1, "sub_vec": [0.1, 0.2, 0.3, 0.4]}]`, schema.GetStructArrayFields()[0])
|
|
require.NoError(t, err)
|
|
fd, err := buildStructArrayFieldData(schema.GetStructArrayFields()[0], []structArrayRow{row})
|
|
require.NoError(t, err)
|
|
|
|
_, err = extractStructArrayRow(fd, 2, schema)
|
|
assert.Error(t, err)
|
|
|
|
mismatch := &schemapb.FieldData{
|
|
Type: schemapb.DataType_ArrayOfStruct,
|
|
FieldName: "my_struct",
|
|
Field: &schemapb.FieldData_StructArrays{
|
|
StructArrays: &schemapb.StructArrayField{Fields: []*schemapb.FieldData{
|
|
{
|
|
Type: schemapb.DataType_Array,
|
|
FieldName: "sub_int",
|
|
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{
|
|
Data: []*schemapb.ScalarField{{
|
|
Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{1, 2}}},
|
|
}},
|
|
}},
|
|
}},
|
|
},
|
|
{
|
|
Type: schemapb.DataType_ArrayOfVector,
|
|
FieldName: "sub_vec",
|
|
Field: &schemapb.FieldData_Vectors{Vectors: &schemapb.VectorField{
|
|
Data: &schemapb.VectorField_VectorArray{VectorArray: &schemapb.VectorArray{
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
Data: []*schemapb.VectorField{{
|
|
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{0.1, 0.2, 0.3, 0.4}}},
|
|
}},
|
|
}},
|
|
}},
|
|
},
|
|
}},
|
|
},
|
|
}
|
|
_, err = extractStructArrayRow(mismatch, 0, schema)
|
|
assert.Error(t, err)
|
|
|
|
missingDimSchema := &schemapb.CollectionSchema{
|
|
StructArrayFields: []*schemapb.StructArrayFieldSchema{{Name: "my_struct"}},
|
|
}
|
|
_, err = extractStructArrayRow(mismatch, 0, missingDimSchema)
|
|
assert.Error(t, err)
|
|
|
|
unsupported := &schemapb.FieldData{
|
|
Type: schemapb.DataType_ArrayOfStruct,
|
|
FieldName: "my_struct",
|
|
Field: &schemapb.FieldData_StructArrays{
|
|
StructArrays: &schemapb.StructArrayField{Fields: []*schemapb.FieldData{
|
|
{
|
|
Type: schemapb.DataType_Array,
|
|
FieldName: "sub_int",
|
|
Field: &schemapb.FieldData_Scalars{Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_ArrayData{ArrayData: &schemapb.ArrayArray{
|
|
Data: []*schemapb.ScalarField{{
|
|
Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{1}}},
|
|
}},
|
|
}},
|
|
}},
|
|
},
|
|
{Type: schemapb.DataType_Bool, FieldName: "bad"},
|
|
}},
|
|
},
|
|
}
|
|
_, err = extractStructArrayRow(unsupported, 0, schema)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestStructArrayHelperValueConversions(t *testing.T) {
|
|
assert.Equal(t, "sub", structFieldShortName("my_struct[sub]"))
|
|
assert.Equal(t, "plain", structFieldShortName("plain"))
|
|
|
|
scalars := []*schemapb.ScalarField{
|
|
{Data: &schemapb.ScalarField_BoolData{BoolData: &schemapb.BoolArray{Data: []bool{true}}}},
|
|
{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{1}}}},
|
|
{Data: &schemapb.ScalarField_LongData{LongData: &schemapb.LongArray{Data: []int64{2}}}},
|
|
{Data: &schemapb.ScalarField_FloatData{FloatData: &schemapb.FloatArray{Data: []float32{3.5}}}},
|
|
{Data: &schemapb.ScalarField_DoubleData{DoubleData: &schemapb.DoubleArray{Data: []float64{4.5}}}},
|
|
{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"five"}}}},
|
|
}
|
|
for _, scalar := range scalars {
|
|
assert.Len(t, scalarArrayToInterfaces(scalar), 1)
|
|
}
|
|
assert.Nil(t, scalarArrayToInterfaces(&schemapb.ScalarField{}))
|
|
|
|
vectorTests := []struct {
|
|
name string
|
|
elementType schemapb.DataType
|
|
vector *schemapb.VectorField
|
|
dim int64
|
|
count int
|
|
}{
|
|
{
|
|
name: "float",
|
|
elementType: schemapb.DataType_FloatVector,
|
|
vector: &schemapb.VectorField{Data: &schemapb.VectorField_FloatVector{
|
|
FloatVector: &schemapb.FloatArray{Data: []float32{0.1, 0.2, 0.3, 0.4}},
|
|
}},
|
|
dim: 2,
|
|
count: 2,
|
|
},
|
|
{
|
|
name: "float16",
|
|
elementType: schemapb.DataType_Float16Vector,
|
|
vector: &schemapb.VectorField{Data: &schemapb.VectorField_Float16Vector{Float16Vector: []byte{1, 2, 3, 4}}},
|
|
dim: 2,
|
|
count: 1,
|
|
},
|
|
{
|
|
name: "bfloat16",
|
|
elementType: schemapb.DataType_BFloat16Vector,
|
|
vector: &schemapb.VectorField{Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: []byte{1, 2, 3, 4}}},
|
|
dim: 2,
|
|
count: 1,
|
|
},
|
|
{
|
|
name: "binary",
|
|
elementType: schemapb.DataType_BinaryVector,
|
|
vector: &schemapb.VectorField{Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{1, 2}}},
|
|
dim: 16,
|
|
count: 1,
|
|
},
|
|
{
|
|
name: "int8",
|
|
elementType: schemapb.DataType_Int8Vector,
|
|
vector: &schemapb.VectorField{Data: &schemapb.VectorField_Int8Vector{Int8Vector: []byte{1, 254}}},
|
|
dim: 2,
|
|
count: 1,
|
|
},
|
|
}
|
|
for _, test := range vectorTests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
count, err := vectorFieldElemCount(test.vector, test.elementType, test.dim)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, test.count, count)
|
|
values, err := vectorFieldToInterfaces(test.vector, test.elementType, test.dim)
|
|
require.NoError(t, err)
|
|
assert.Len(t, values, test.count)
|
|
})
|
|
}
|
|
|
|
_, err := vectorFieldElemCount(&schemapb.VectorField{}, schemapb.DataType_FloatVector, 0)
|
|
assert.Error(t, err)
|
|
_, err = vectorFieldElemCount(&schemapb.VectorField{}, schemapb.DataType_JSON, 1)
|
|
assert.Error(t, err)
|
|
_, err = vectorFieldToInterfaces(&schemapb.VectorField{}, schemapb.DataType_FloatVector, 0)
|
|
assert.Error(t, err)
|
|
_, err = vectorFieldToInterfaces(&schemapb.VectorField{}, schemapb.DataType_JSON, 1)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestStructArrayCheckAndSetPartialUpdate(t *testing.T) {
|
|
schema := buildStructArrayTestSchema()
|
|
body := []byte(`{"data": [{"id": 1, "vec": [0.1,0.2,0.3,0.4]}]}`)
|
|
rows, validData, err := checkAndSetData(body, schema, true)
|
|
require.NoError(t, err)
|
|
assert.NotContains(t, rows[0], "my_struct")
|
|
fds, err := anyToColumns(rows, validData, schema, false, true)
|
|
require.NoError(t, err)
|
|
for _, fd := range fds {
|
|
assert.NotEqual(t, schemapb.DataType_ArrayOfStruct, fd.GetType())
|
|
}
|
|
|
|
_, _, err = checkAndSetData(body, schema, false)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestStructArrayFieldDataValueCount(t *testing.T) {
|
|
schema := buildStructArrayTestSchema().GetStructArrayFields()[0]
|
|
row, err := parseStructArrayRow(`[{"sub_int": 1, "sub_vec": [0.1,0.2,0.3,0.4]}]`, schema)
|
|
require.NoError(t, err)
|
|
fd, err := buildStructArrayFieldData(schema, []structArrayRow{row})
|
|
require.NoError(t, err)
|
|
count, err := fieldDataValueCount(fd)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(1), count)
|
|
|
|
subs := fd.GetStructArrays().GetFields()
|
|
fd.GetStructArrays().Fields = []*schemapb.FieldData{subs[1], subs[0]}
|
|
count, err = fieldDataValueCount(fd)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(1), count)
|
|
|
|
fd.GetStructArrays().Fields = nil
|
|
count, err = fieldDataValueCount(fd)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(0), count)
|
|
|
|
fd.GetStructArrays().Fields = []*schemapb.FieldData{{Type: schemapb.DataType_Bool}}
|
|
_, err = fieldDataValueCount(fd)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestStructArrayFieldSchemaGetProtoTypeParams(t *testing.T) {
|
|
proto, err := (&StructArrayFieldSchema{
|
|
FieldName: "my_struct",
|
|
Description: "with params",
|
|
Nullable: true,
|
|
TypeParams: map[string]interface{}{
|
|
common.MaxCapacityKey: 8,
|
|
},
|
|
Fields: []FieldSchema{
|
|
{FieldName: "sub_int", DataType: "Array", ElementDataType: "Int32"},
|
|
},
|
|
}).GetProto(context.Background())
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "my_struct", proto.GetName())
|
|
assert.Equal(t, "with params", proto.GetDescription())
|
|
assert.True(t, proto.GetNullable())
|
|
require.Len(t, proto.GetTypeParams(), 1)
|
|
assert.Equal(t, common.MaxCapacityKey, proto.GetTypeParams()[0].GetKey())
|
|
assert.Equal(t, "8", proto.GetTypeParams()[0].GetValue())
|
|
subParams := proto.GetFields()[0].GetTypeParams()
|
|
require.Len(t, subParams, 1)
|
|
assert.Equal(t, common.MaxCapacityKey, subParams[0].GetKey())
|
|
assert.Equal(t, "8", subParams[0].GetValue())
|
|
|
|
_, err = (&StructArrayFieldSchema{
|
|
FieldName: "bad_params",
|
|
TypeParams: map[string]interface{}{
|
|
"bad": func() {},
|
|
},
|
|
Fields: []FieldSchema{
|
|
{FieldName: "sub_int", DataType: "Array", ElementDataType: "Int32"},
|
|
},
|
|
}).GetProto(context.Background())
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestFieldSchemaStructArrayHelpers(t *testing.T) {
|
|
var nilField *FieldSchema
|
|
assert.False(t, nilField.IsStructArrayField())
|
|
assert.False(t, (&FieldSchema{DataType: "Int64"}).IsStructArrayField())
|
|
assert.True(t, (&FieldSchema{DataType: "Array", ElementDataType: "Struct"}).IsStructArrayField())
|
|
assert.True(t, (&FieldSchema{DataType: "ArrayOfStruct"}).IsStructArrayField())
|
|
|
|
proto, err := (&FieldSchema{
|
|
FieldName: "clips",
|
|
Description: "clip metadata",
|
|
DataType: "Array",
|
|
ElementDataType: "Struct",
|
|
Nullable: true,
|
|
ElementTypeParams: map[string]interface{}{
|
|
common.MaxCapacityKey: 16,
|
|
},
|
|
TypeParams: map[string]interface{}{
|
|
common.MaxCapacityKey: 32,
|
|
},
|
|
Fields: []FieldSchema{
|
|
{
|
|
FieldName: "tag",
|
|
DataType: "Array",
|
|
ElementDataType: "VarChar",
|
|
ElementTypeParams: map[string]interface{}{
|
|
common.MaxLengthKey: 64,
|
|
},
|
|
},
|
|
{
|
|
FieldName: "scores",
|
|
DataType: "Array",
|
|
ElementDataType: "Int64",
|
|
ElementTypeParams: map[string]interface{}{
|
|
common.MaxCapacityKey: 7,
|
|
},
|
|
},
|
|
},
|
|
}).GetStructArrayProto(context.Background())
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "clips", proto.GetName())
|
|
assert.Equal(t, "clip metadata", proto.GetDescription())
|
|
assert.True(t, proto.GetNullable())
|
|
assert.Equal(t, "32", kvPairsToMap(proto.GetTypeParams())[common.MaxCapacityKey])
|
|
|
|
require.Len(t, proto.GetFields(), 2)
|
|
tagParams := kvPairsToMap(proto.GetFields()[0].GetTypeParams())
|
|
assert.Equal(t, "64", tagParams[common.MaxLengthKey])
|
|
assert.Equal(t, "32", tagParams[common.MaxCapacityKey])
|
|
|
|
scoreParams := kvPairsToMap(proto.GetFields()[1].GetTypeParams())
|
|
assert.Equal(t, "7", scoreParams[common.MaxCapacityKey])
|
|
}
|
|
|
|
func TestPrintStructArrayFieldsV2QualifiedSubFields(t *testing.T) {
|
|
printed := printStructArrayFieldsV2([]*schemapb.StructArrayFieldSchema{
|
|
{
|
|
FieldID: 10,
|
|
Name: "my_struct",
|
|
Description: "qualified names",
|
|
TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxCapacityKey, Value: "16"}},
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: 11,
|
|
Name: "my_struct[sub_int]",
|
|
DataType: schemapb.DataType_Array,
|
|
ElementType: schemapb.DataType_Int32,
|
|
},
|
|
{
|
|
FieldID: 12,
|
|
Name: "my_struct[sub_vec]",
|
|
DataType: schemapb.DataType_ArrayOfVector,
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
require.Len(t, printed, 1)
|
|
entry := printed[0]
|
|
assert.Equal(t, "my_struct", entry[HTTPReturnFieldName])
|
|
params := entry[Params].([]*commonpb.KeyValuePair)
|
|
require.Len(t, params, 1)
|
|
assert.Equal(t, common.MaxCapacityKey, params[0].GetKey())
|
|
assert.Equal(t, "16", params[0].GetValue())
|
|
|
|
subs := entry["fields"].([]gin.H)
|
|
require.Len(t, subs, 2)
|
|
assert.Equal(t, "sub_int", subs[0][HTTPReturnFieldName])
|
|
assert.Equal(t, schemapb.DataType_Int32.String(), subs[0][HTTPReturnFieldElementType])
|
|
assert.Equal(t, "sub_vec", subs[1][HTTPReturnFieldName])
|
|
assert.Equal(t, schemapb.DataType_FloatVector.String(), subs[1][HTTPReturnFieldElementType])
|
|
}
|
|
|
|
func TestCheckAndSetDataStructArrayRows(t *testing.T) {
|
|
schema := buildStructArrayTestSchema()
|
|
body := []byte(`{"data": [
|
|
{
|
|
"id": 1,
|
|
"vec": [0.1, 0.2, 0.3, 0.4],
|
|
"my_struct": [
|
|
{"sub_int": 10, "sub_vec": [1.1, 1.2, 1.3, 1.4], "ignored": true}
|
|
]
|
|
}
|
|
]}`)
|
|
|
|
rows, validData, err := checkAndSetData(body, schema, false)
|
|
require.NoError(t, err)
|
|
require.Empty(t, validData)
|
|
require.Len(t, rows, 1)
|
|
structRow, ok := rows[0]["my_struct"].(structArrayRow)
|
|
require.True(t, ok)
|
|
assert.Contains(t, structRow, "sub_int")
|
|
assert.Contains(t, structRow, "sub_vec")
|
|
|
|
_, _, err = checkAndSetData([]byte(`{"data": [{"id": 1, "vec": [0.1,0.2,0.3,0.4], "my_struct": [1]}]}`), schema, false)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestParseStructArrayRowQualifiedSchemaNames(t *testing.T) {
|
|
schema := &schemapb.StructArrayFieldSchema{
|
|
Name: "my_struct",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
Name: "my_struct[sub_int]",
|
|
DataType: schemapb.DataType_Array,
|
|
ElementType: schemapb.DataType_Int32,
|
|
},
|
|
{
|
|
Name: "my_struct[sub_vec]",
|
|
DataType: schemapb.DataType_ArrayOfVector,
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "2"}},
|
|
},
|
|
},
|
|
}
|
|
row, err := parseStructArrayRow(`[{"sub_int": 1, "sub_vec": [0.1,0.2], "unknown": "ignored"}]`, schema)
|
|
require.NoError(t, err)
|
|
assert.Contains(t, row, "sub_int")
|
|
assert.Contains(t, row, "sub_vec")
|
|
assert.NotContains(t, row, "unknown")
|
|
|
|
unsupported := proto.Clone(schema).(*schemapb.StructArrayFieldSchema)
|
|
unsupported.Fields = append(unsupported.Fields, &schemapb.FieldSchema{Name: "bad", DataType: schemapb.DataType_Bool})
|
|
_, err = parseStructArrayRow(`[{"sub_int": 1, "sub_vec": [0.1,0.2], "bad": true}]`, unsupported)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestByteVectorElementErrorPaths(t *testing.T) {
|
|
_, err := decodeByteVectorElement(gjson.Parse(`"bad-base64"`), 2, 4, true)
|
|
assert.Error(t, err)
|
|
|
|
_, err = decodeByteVectorElement(gjson.Parse(`"AQI="`), 2, 4, true)
|
|
assert.Error(t, err)
|
|
|
|
_, err = decodeByteVectorElement(gjson.Parse(`true`), 2, 4, true)
|
|
assert.Error(t, err)
|
|
|
|
_, err = decodeByteVectorElement(gjson.Parse(`[0.1]`), 2, 4, true)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestEncodeEmbListQueryErrorPaths(t *testing.T) {
|
|
_, err := encodeEmbListQuery(gjson.Parse(`[[0.1]]`).Array(), schemapb.DataType_FloatVector, 2, 0)
|
|
assert.Error(t, err)
|
|
|
|
_, err = encodeEmbListQuery(gjson.Parse(`["AQI="]`).Array(), schemapb.DataType_BinaryVector, 32, 0)
|
|
assert.Error(t, err)
|
|
|
|
_, err = encodeEmbListQuery(gjson.Parse(`[[1, 2, 300]]`).Array(), schemapb.DataType_Int8Vector, 3, 0)
|
|
assert.Error(t, err)
|
|
|
|
_, err = encodeEmbListQuery(gjson.Parse(`[[true]]`).Array(), schemapb.DataType_Bool, 1, 0)
|
|
assert.Error(t, err)
|
|
}
|