diff --git a/internal/querynode/mock_test.go b/internal/querynode/mock_test.go index 51d863a753..fbda74c23a 100644 --- a/internal/querynode/mock_test.go +++ b/internal/querynode/mock_test.go @@ -1966,6 +1966,36 @@ func genFieldData(fieldName string, fieldID int64, fieldType schemapb.DataType, }, FieldId: fieldID, } + case schemapb.DataType_JSON: + fieldData = &schemapb.FieldData{ + Type: schemapb.DataType_JSON, + FieldName: fieldName, + Field: &schemapb.FieldData_Scalars{ + Scalars: &schemapb.ScalarField{ + Data: &schemapb.ScalarField_JsonData{ + JsonData: &schemapb.JSONArray{ + Data: fieldValue.([][]byte), + }, + }, + }, + }, + FieldId: fieldID, + } + case schemapb.DataType_Array: + fieldData = &schemapb.FieldData{ + Type: schemapb.DataType_Array, + FieldName: fieldName, + Field: &schemapb.FieldData_Scalars{ + Scalars: &schemapb.ScalarField{ + Data: &schemapb.ScalarField_ArrayData{ + ArrayData: &schemapb.ArrayArray{ + Data: fieldValue.([]*schemapb.ScalarField), + }, + }, + }, + }, + FieldId: fieldID, + } default: log.Error("not supported field type", zap.String("field type", fieldType.String())) } diff --git a/internal/querynode/result_sorter.go b/internal/querynode/result_sorter.go index 526c1a52a7..f219e2753b 100644 --- a/internal/querynode/result_sorter.go +++ b/internal/querynode/result_sorter.go @@ -2,6 +2,7 @@ package querynode import ( "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" + "github.com/milvus-io/milvus/internal/log" "github.com/milvus-io/milvus/internal/proto/segcorepb" "github.com/milvus-io/milvus/internal/util/typeutil" ) @@ -60,6 +61,16 @@ func swapFieldData(field *schemapb.FieldData, i int, j int) { case *schemapb.ScalarField_StringData: data := sd.StringData.Data data[i], data[j] = data[j], data[i] + case *schemapb.ScalarField_JsonData: + data := sd.JsonData.Data + data[i], data[j] = data[j], data[i] + case *schemapb.ScalarField_ArrayData: + data := sd.ArrayData.Data + data[i], data[j] = data[j], data[i] + default: + errMsg := "undefined data type " + field.Type.String() + log.Error(errMsg) + panic(errMsg) } case *schemapb.FieldData_Vectors: dim := int(field.GetVectors().GetDim()) diff --git a/internal/querynode/result_sorter_test.go b/internal/querynode/result_sorter_test.go index 737bbcca1a..af726a3d0b 100644 --- a/internal/querynode/result_sorter_test.go +++ b/internal/querynode/result_sorter_test.go @@ -36,6 +36,20 @@ func TestResultSorter_ByIntPK(t *testing.T) { []float32{5, 4, 3, 2, 9, 8, 7, 6}, 1), genFieldData("binary vector field", 107, schemapb.DataType_BinaryVector, []byte{5, 4, 3, 2, 9, 8, 7, 6}, 8), + genFieldData("json field", 108, schemapb.DataType_JSON, + [][]byte{[]byte("{\"5\": 5}"), []byte("{\"4\": 4}"), []byte("{\"3\": 3}"), []byte("{\"2\": 2}"), + []byte("{\"9\": 9}"), []byte("{\"8\": 8}"), []byte("{\"7\": 7}"), []byte("{\"6\": 6}")}, 1), + genFieldData("json field", 108, schemapb.DataType_Array, + []*schemapb.ScalarField{ + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{5, 6, 7}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{4, 5, 6}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{3, 4, 5}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{2, 3, 4}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{9, 10, 11}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{8, 9, 10}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{7, 8, 9}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{6, 7, 8}}}}, + }, 1), }, } @@ -51,4 +65,16 @@ func TestResultSorter_ByIntPK(t *testing.T) { assert.Equal(t, []int32{2, 3, 4, 5, 6, 7, 8, 9}, result.FieldsData[5].GetScalars().GetIntData().Data) assert.InDeltaSlice(t, []float32{2, 3, 4, 5, 6, 7, 8, 9}, result.FieldsData[6].GetVectors().GetFloatVector().GetData(), 10e-10) assert.Equal(t, []byte{2, 3, 4, 5, 6, 7, 8, 9}, result.FieldsData[7].GetVectors().GetBinaryVector()) + assert.Equal(t, [][]byte{[]byte("{\"2\": 2}"), []byte("{\"3\": 3}"), []byte("{\"4\": 4}"), []byte("{\"5\": 5}"), + []byte("{\"6\": 6}"), []byte("{\"7\": 7}"), []byte("{\"8\": 8}"), []byte("{\"9\": 9}")}, result.FieldsData[8].GetScalars().GetJsonData().GetData()) + assert.Equal(t, []*schemapb.ScalarField{ + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{2, 3, 4}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{3, 4, 5}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{4, 5, 6}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{5, 6, 7}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{6, 7, 8}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{7, 8, 9}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{8, 9, 10}}}}, + {Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{9, 10, 11}}}}, + }, result.FieldsData[9].GetScalars().GetArrayData().GetData()) }