Add wildcard outputfield expansion logic back

Signed-off-by: Congqi Xia <[email protected]>

client/data.go

@@ -92,7 +92,8 @@ func (c *GrpcClient) Search(ctx context.Context, collName string, partitions []s
 	return sr, nil
 }
 
-func (c *GrpcClient) parseSearchResult(_ *entity.Schema, outputFields []string, fieldDataList []*schemapb.FieldData, _, from, to int) ([]entity.Column, error) {
+func (c *GrpcClient) parseSearchResult(sch *entity.Schema, outputFields []string, fieldDataList []*schemapb.FieldData, _, from, to int) ([]entity.Column, error) {
+	outputFields = expandWildcard(sch, outputFields)
 	// duplicated name will have only one column now
 	outputSet := make(map[string]struct{})
 	for _, output := range outputFields {
@@ -141,6 +142,37 @@ func (c *GrpcClient) parseSearchResult(_ *entity.Schema, outputFields []string,
 	return columns, nil
 }
 
+func expandWildcard(schema *entity.Schema, outputFields []string) []string {
+	wildcard := false
+	for _, outputField := range outputFields {
+		if outputField == "*" {
+			wildcard = true
+		}
+	}
+	if !wildcard {
+		return outputFields
+	}
+
+	set := make(map[string]struct{})
+	result := make([]string, 0, len(schema.Fields))
+	for _, field := range schema.Fields {
+		result = append(result, field.Name)
+		set[field.Name] = struct{}{}
+	}
+
+	// add dynamic fields output
+	for _, output := range outputFields {
+		if output == "*" {
+			continue
+		}
+		_, ok := set[output]
+		if !ok {
+			result = append(result, output)
+		}
+	}
+	return result
+}
+
 func PKs2Expr(backName string, ids entity.Column) string {
 	var expr string
 	var pkName = ids.Name()

client/data_test.go

@@ -1178,3 +1178,42 @@ func TestVector2PlaceHolder(t *testing.T) {
 		}
 	})
 }
+
+type WildcardSuite struct {
+	suite.Suite
+
+	schema *entity.Schema
+}
+
+func (s *WildcardSuite) SetupTest() {
+	s.schema = entity.NewSchema().
+		WithField(entity.NewField().WithName("pk").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
+		WithField(entity.NewField().WithName("attr").WithDataType(entity.FieldTypeInt64)).
+		WithField(entity.NewField().WithName("$meta").WithDataType(entity.FieldTypeJSON).WithIsDynamic(true)).
+		WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(128))
+}
+
+func (s *WildcardSuite) TestExpandWildcard() {
+	type testCase struct {
+		tag    string
+		input  []string
+		expect []string
+	}
+
+	cases := []testCase{
+		{tag: "normal", input: []string{"pk", "attr"}, expect: []string{"pk", "attr"}},
+		{tag: "with_wildcard", input: []string{"*"}, expect: []string{"pk", "attr", "$meta", "vector"}},
+		{tag: "wildcard_dynamic", input: []string{"*", "a"}, expect: []string{"pk", "attr", "$meta", "vector", "a"}},
+	}
+
+	for _, tc := range cases {
+		s.Run(tc.tag, func() {
+			output := expandWildcard(s.schema, tc.input)
+			s.ElementsMatch(tc.expect, output)
+		})
+	}
+}
+
+func TestExpandWildcard(t *testing.T) {
+	suite.Run(t, new(WildcardSuite))
+}