Handle LongData returned when queryresult is empty

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
milvus-io · May 4, 2023 · c6d22de · c6d22de
1 parent 0f8c85b
commit c6d22de
Show file tree

Hide file tree

Showing 2 changed files with 97 additions and 3 deletions.
diff --git a/entity/columns.go b/entity/columns.go
@@ -124,7 +124,7 @@ func FieldDataColumn(fd *schema.FieldData, begin, end int) (Column, error) {
 		return NewColumnBool(fd.GetFieldName(), data.BoolData.GetData()[begin:end]), nil
 
 	case schema.DataType_Int8:
-		data, ok := fd.GetScalars().GetData().(*schema.ScalarField_IntData)
+		data, ok := getIntData(fd)
 		if !ok {
 			return nil, errFieldDataTypeNotMatch
 		}
@@ -140,7 +140,7 @@ func FieldDataColumn(fd *schema.FieldData, begin, end int) (Column, error) {
 		return NewColumnInt8(fd.GetFieldName(), values[begin:end]), nil
 
 	case schema.DataType_Int16:
-		data, ok := fd.GetScalars().GetData().(*schema.ScalarField_IntData)
+		data, ok := getIntData(fd)
 		if !ok {
 			return nil, errFieldDataTypeNotMatch
 		}
@@ -155,7 +155,7 @@ func FieldDataColumn(fd *schema.FieldData, begin, end int) (Column, error) {
 		return NewColumnInt16(fd.GetFieldName(), values[begin:end]), nil
 
 	case schema.DataType_Int32:
-		data, ok := fd.GetScalars().GetData().(*schema.ScalarField_IntData)
+		data, ok := getIntData(fd)
 		if !ok {
 			return nil, errFieldDataTypeNotMatch
 		}
@@ -218,6 +218,25 @@ func FieldDataColumn(fd *schema.FieldData, begin, end int) (Column, error) {
 	}
 }
 
+// getIntData get int32 slice from result field data
+// also handles LongData bug (see also https://github.com/milvus-io/milvus/issues/23850)
+func getIntData(fd *schema.FieldData) (*schema.ScalarField_IntData, bool) {
+	switch data := fd.GetScalars().GetData().(type) {
+	case *schema.ScalarField_IntData:
+		return data, true
+	case *schema.ScalarField_LongData:
+		// only alway empty LongData for backward compatibility
+		if len(data.LongData.GetData()) == 0 {
+			return &schema.ScalarField_IntData{
+				IntData: &schema.IntArray{},
+			}, true
+		}
+		return nil, false
+	default:
+		return nil, false
+	}
+}
+
 // FieldDataColumn converts schema.FieldData to vector Column
 func FieldDataVector(fd *schema.FieldData) (Column, error) {
 	switch fd.GetType() {

diff --git a/entity/columns_test.go b/entity/columns_test.go
@@ -95,3 +95,78 @@ func TestIDColumns(t *testing.T) {
 		assert.Equal(t, dataLen, column.Len())
 	})
 }
+
+func TestGetIntData(t *testing.T) {
+	type testCase struct {
+		tag      string
+		fd       *schema.FieldData
+		expectOK bool
+	}
+
+	cases := []testCase{
+		{
+			tag: "normal_IntData",
+			fd: &schema.FieldData{
+				Field: &schema.FieldData_Scalars{
+					Scalars: &schema.ScalarField{
+						Data: &schema.ScalarField_IntData{
+							IntData: &schema.IntArray{Data: []int32{1, 2, 3}},
+						},
+					},
+				},
+			},
+			expectOK: true,
+		},
+		{
+			tag: "empty_LongData",
+			fd: &schema.FieldData{
+				Field: &schema.FieldData_Scalars{
+					Scalars: &schema.ScalarField{
+						Data: &schema.ScalarField_LongData{
+							LongData: &schema.LongArray{Data: nil},
+						},
+					},
+				},
+			},
+			expectOK: true,
+		},
+		{
+			tag: "nonempty_LongData",
+			fd: &schema.FieldData{
+				Field: &schema.FieldData_Scalars{
+					Scalars: &schema.ScalarField{
+						Data: &schema.ScalarField_LongData{
+							LongData: &schema.LongArray{Data: []int64{1, 2, 3}},
+						},
+					},
+				},
+			},
+			expectOK: false,
+		},
+		{
+			tag: "other_data",
+			fd: &schema.FieldData{
+				Field: &schema.FieldData_Scalars{
+					Scalars: &schema.ScalarField{
+						Data: &schema.ScalarField_BoolData{},
+					},
+				},
+			},
+			expectOK: false,
+		},
+		{
+			tag: "vector_data",
+			fd: &schema.FieldData{
+				Field: &schema.FieldData_Vectors{},
+			},
+			expectOK: false,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.tag, func(t *testing.T) {
+			_, ok := getIntData(tc.fd)
+			assert.Equal(t, tc.expectOK, ok)
+		})
+	}
+}