enhance: Support dynamic field in SchemaHelper (#35461)

Related to #35415

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2024-08-14 18:04:59 +08:00 committed by GitHub
parent 0400eb91c1
commit f99a761d87
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 136 additions and 2 deletions

View File

@ -37,7 +37,6 @@ func newTestSchema() *schemapb.CollectionSchema {
FieldID: 131, Name: "StringArrayField", IsPrimaryKey: false, Description: "string array field",
DataType: schemapb.DataType_Array,
ElementType: schemapb.DataType_VarChar,
IsDynamic: true,
})
return &schemapb.CollectionSchema{

View File

@ -256,6 +256,7 @@ type SchemaHelper struct {
idOffset map[int64]int
primaryKeyOffset int
partitionKeyOffset int
dynamicFieldOffset int
}
// CreateSchemaHelper returns a new SchemaHelper object
@ -263,7 +264,7 @@ func CreateSchemaHelper(schema *schemapb.CollectionSchema) (*SchemaHelper, error
if schema == nil {
return nil, errors.New("schema is nil")
}
schemaHelper := SchemaHelper{schema: schema, nameOffset: make(map[string]int), idOffset: make(map[int64]int), primaryKeyOffset: -1, partitionKeyOffset: -1}
schemaHelper := SchemaHelper{schema: schema, nameOffset: make(map[string]int), idOffset: make(map[int64]int), primaryKeyOffset: -1, partitionKeyOffset: -1, dynamicFieldOffset: -1}
for offset, field := range schema.Fields {
if _, ok := schemaHelper.nameOffset[field.Name]; ok {
return nil, fmt.Errorf("duplicated fieldName: %s", field.Name)
@ -286,6 +287,13 @@ func CreateSchemaHelper(schema *schemapb.CollectionSchema) (*SchemaHelper, error
}
schemaHelper.partitionKeyOffset = offset
}
if field.IsDynamic {
if schemaHelper.dynamicFieldOffset != -1 {
return nil, errors.New("dynamic field is not unique")
}
schemaHelper.dynamicFieldOffset = offset
}
}
return &schemaHelper, nil
}
@ -306,6 +314,15 @@ func (helper *SchemaHelper) GetPartitionKeyField() (*schemapb.FieldSchema, error
return helper.schema.Fields[helper.partitionKeyOffset], nil
}
// GetDynamicField returns the field schema of dynamic field if exists.
// if there is no dynamic field defined in schema, error will be returned.
func (helper *SchemaHelper) GetDynamicField() (*schemapb.FieldSchema, error) {
if helper.dynamicFieldOffset == -1 {
return nil, fmt.Errorf("failed to get dynamic field: no dynamic field in schema")
}
return helper.schema.Fields[helper.dynamicFieldOffset], nil
}
// GetFieldFromName is used to find the schema by field name
func (helper *SchemaHelper) GetFieldFromName(fieldName string) (*schemapb.FieldSchema, error) {
offset, ok := helper.nameOffset[fieldName]

View File

@ -356,6 +356,124 @@ func TestSchema_GetVectorFieldSchema(t *testing.T) {
})
}
func TestSchemaHelper_GetDynamicField(t *testing.T) {
t.Run("with_dynamic_schema", func(t *testing.T) {
sch := &schemapb.CollectionSchema{
Name: "testColl",
Description: "",
AutoID: false,
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "field_int64",
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
},
{
FieldID: 101,
Name: "field_float_vector",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
{
FieldID: 102,
Name: "$meta",
DataType: schemapb.DataType_JSON,
IsDynamic: true,
},
},
}
helper, err := CreateSchemaHelper(sch)
require.NoError(t, err)
f, err := helper.GetDynamicField()
assert.NoError(t, err)
assert.NotNil(t, f)
assert.EqualValues(t, 102, f.FieldID)
})
t.Run("without_dynamic_schema", func(t *testing.T) {
sch := &schemapb.CollectionSchema{
Name: "testColl",
Description: "",
AutoID: false,
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "field_int64",
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
},
{
FieldID: 101,
Name: "field_float_vector",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
},
}
helper, err := CreateSchemaHelper(sch)
require.NoError(t, err)
_, err = helper.GetDynamicField()
assert.Error(t, err)
})
t.Run("multiple_dynamic_fields", func(t *testing.T) {
sch := &schemapb.CollectionSchema{
Name: "testColl",
Description: "",
AutoID: false,
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "field_int64",
IsPrimaryKey: true,
DataType: schemapb.DataType_Int64,
},
{
FieldID: 101,
Name: "field_float_vector",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
},
{
FieldID: 102,
Name: "$meta",
DataType: schemapb.DataType_JSON,
IsDynamic: true,
},
{
FieldID: 103,
Name: "other_json",
DataType: schemapb.DataType_JSON,
IsDynamic: true,
},
},
}
_, err := CreateSchemaHelper(sch)
assert.Error(t, err)
})
}
func TestSchema_invalid(t *testing.T) {
t.Run("Duplicate field name", func(t *testing.T) {
schema := &schemapb.CollectionSchema{