Support default value

Signed-off-by: lixinguo <xinguo.li@zilliz.com>
milvus-io · Jul 16, 2023 · 116a963 · 116a963
1 parent 58e8817
commit 116a963
Show file tree

Hide file tree

Showing 5 changed files with 123 additions and 2 deletions.
diff --git a/client/insert.go b/client/insert.go
@@ -140,6 +140,15 @@ func (c *GrpcClient) processInsertColumns(colSchema *entity.Schema, columns ...e
 		_, has := mNameColumn[field.Name]
 		if !has &&
 			!field.AutoID && !field.IsDynamic {
+			if field.DefaultValue != nil {
+				emptyColumn, err := entity.DefaultValueColumn(field.Name, field.DataType)
+				// no err will be throw
+				if err != nil {
+					return nil, 0, fmt.Errorf("empty column fail to create, field name:%s ", field.Name)
+				}
+				mNameColumn[field.Name] = emptyColumn
+				continue
+			}
 			return nil, 0, fmt.Errorf("field %s not passed", field.Name)
 		}
 	}

diff --git a/client/insert_test.go b/client/insert_test.go
@@ -64,7 +64,7 @@ func (s *InsertSuite) TestInsertFail() {
 		s.Error(err)
 	})
 
-	s.Run("missing_field", func() {
+	s.Run("missing_field_without_default_value", func() {
 		defer s.resetMock()
 		s.setupHasCollection(testCollectionName)
 		s.setupHasPartition(testCollectionName, "partition_1")
@@ -80,6 +80,24 @@ func (s *InsertSuite) TestInsertFail() {
 		s.Error(err)
 	})
 
+	s.Run("missing_field_with_default_value", func() {
+		defer s.resetMock()
+		s.setupHasCollection(testCollectionName)
+		s.setupHasPartition(testCollectionName, "partition_1")
+
+		s.setupDescribeCollection(testCollectionName, entity.NewSchema().
+			WithField(entity.NewField().WithIsPrimaryKey(true).WithIsAutoID(true).WithName("ID").WithDataType(entity.FieldTypeInt64)).
+			WithField(entity.NewField().WithName("default_value").WithDataType(entity.FieldTypeInt64).WithDefaultValueLong(1)).
+			WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithTypeParams(entity.TypeParamDim, "128")),
+		)
+
+		_, err := c.Insert(ctx, testCollectionName, "partition_1",
+			entity.NewColumnInt64("ID", []int64{1}),
+			entity.NewColumnFloatVector("vector", 128, [][]float32{}),
+		)
+		s.NoError(err)
+	})
+
 	s.Run("column_len_not_match", func() {
 		defer s.resetMock()
 		s.setupHasCollection(testCollectionName)
@@ -92,7 +110,7 @@ func (s *InsertSuite) TestInsertFail() {
 
 		_, err := c.Insert(ctx, testCollectionName, "partition_1",
 			entity.NewColumnInt64("ID", []int64{1, 2}),
-			entity.NewColumnFloatVector("vector", 128, [][]float32{}),
+			entity.NewColumnFloatVector("vector", 128, generateFloatVector(1, 128)),
 		)
 		s.Error(err)
 	})

diff --git a/entity/columns.go b/entity/columns.go
@@ -357,3 +357,32 @@ func FieldDataVector(fd *schema.FieldData) (Column, error) {
 		return nil, errors.New("unsupported data type")
 	}
 }
+
+// defaultValueColumn will return the empty scalars column which will be fill with default value
+func DefaultValueColumn(name string, dataType FieldType) (Column, error) {
+	switch dataType {
+	case FieldTypeBool:
+		return NewColumnBool(name, nil), nil
+	case FieldTypeInt8:
+		return NewColumnInt8(name, nil), nil
+	case FieldTypeInt16:
+		return NewColumnInt16(name, nil), nil
+	case FieldTypeInt32:
+		return NewColumnInt32(name, nil), nil
+	case FieldTypeInt64:
+		return NewColumnInt64(name, nil), nil
+	case FieldTypeFloat:
+		return NewColumnFloat(name, nil), nil
+	case FieldTypeDouble:
+		return NewColumnDouble(name, nil), nil
+	case FieldTypeString:
+		return NewColumnString(name, nil), nil
+	case FieldTypeVarChar:
+		return NewColumnVarChar(name, nil), nil
+	case FieldTypeJSON:
+		return NewColumnJSONBytes(name, nil), nil
+
+	default:
+		return nil, fmt.Errorf("default value unsupported data type %s", dataType)
+	}
+}
diff --git a/entity/schema.go b/entity/schema.go
@@ -137,6 +137,7 @@ type Field struct {
 	IndexParams    map[string]string
 	IsDynamic      bool
 	IsPartitionKey bool
+	DefaultValue   *schema.ValueField
 }
 
 // ProtoMessage generates corresponding FieldSchema
@@ -152,6 +153,7 @@ func (f *Field) ProtoMessage() *schema.FieldSchema {
 		IndexParams:    MapKvPairs(f.IndexParams),
 		IsDynamic:      f.IsDynamic,
 		IsPartitionKey: f.IsPartitionKey,
+		DefaultValue:   f.DefaultValue,
 	}
 }
 
@@ -198,6 +200,60 @@ func (f *Field) WithIsPartitionKey(isPartitionKey bool) *Field {
 	return f
 }
 
+func (f *Field) WithDefaultValueBool(defaultValue bool) *Field {
+	f.DefaultValue = &schema.ValueField{
+		Data: &schema.ValueField_BoolData{
+			BoolData: defaultValue,
+		},
+	}
+	return f
+}
+
+func (f *Field) WithDefaultValueInt(defaultValue int32) *Field {
+	f.DefaultValue = &schema.ValueField{
+		Data: &schema.ValueField_IntData{
+			IntData: defaultValue,
+		},
+	}
+	return f
+}
+
+func (f *Field) WithDefaultValueLong(defaultValue int64) *Field {
+	f.DefaultValue = &schema.ValueField{
+		Data: &schema.ValueField_LongData{
+			LongData: defaultValue,
+		},
+	}
+	return f
+}
+
+func (f *Field) WithDefaultValueFloat(defaultValue float32) *Field {
+	f.DefaultValue = &schema.ValueField{
+		Data: &schema.ValueField_FloatData{
+			FloatData: defaultValue,
+		},
+	}
+	return f
+}
+
+func (f *Field) WithDefaultValueDouble(defaultValue float64) *Field {
+	f.DefaultValue = &schema.ValueField{
+		Data: &schema.ValueField_DoubleData{
+			DoubleData: defaultValue,
+		},
+	}
+	return f
+}
+
+func (f *Field) WithDefaultValueString(defaultValue string) *Field {
+	f.DefaultValue = &schema.ValueField{
+		Data: &schema.ValueField_StringData{
+			StringData: defaultValue,
+		},
+	}
+	return f
+}
+
 func (f *Field) WithTypeParams(key string, value string) *Field {
 	if f.TypeParams == nil {
 		f.TypeParams = make(map[string]string)
@@ -234,6 +290,7 @@ func (f *Field) ReadProto(p *schema.FieldSchema) *Field {
 	f.IndexParams = KvPairsMap(p.GetIndexParams())
 	f.IsDynamic = p.GetIsDynamic()
 	f.IsPartitionKey = p.GetIsPartitionKey()
+	f.DefaultValue = p.GetDefaultValue()
 
 	return f
 }

diff --git a/entity/schema_test.go b/entity/schema_test.go
@@ -25,6 +25,12 @@ func TestFieldSchema(t *testing.T) {
 		NewField().WithName("int_field").WithDataType(FieldTypeInt64).WithIsAutoID(true).WithIsPrimaryKey(true).WithDescription("int_field desc"),
 		NewField().WithName("string_field").WithDataType(FieldTypeString).WithIsAutoID(false).WithIsPrimaryKey(true).WithIsDynamic(false).WithTypeParams("max_len", "32").WithDescription("string_field desc"),
 		NewField().WithName("partition_key").WithDataType(FieldTypeInt32).WithIsPartitionKey(true),
+		NewField().WithName("default_value_bool").WithDataType(FieldTypeBool).WithDefaultValueBool(true),
+		NewField().WithName("default_value_int").WithDataType(FieldTypeInt32).WithDefaultValueInt(1),
+		NewField().WithName("default_value_long").WithDataType(FieldTypeInt64).WithDefaultValueLong(1),
+		NewField().WithName("default_value_float").WithDataType(FieldTypeFloat).WithDefaultValueFloat(1),
+		NewField().WithName("default_value_double").WithDataType(FieldTypeDouble).WithDefaultValueDouble(1),
+		NewField().WithName("default_value_string").WithDataType(FieldTypeString).WithDefaultValueString("a"),
 	}
 
 	for _, field := range fields {
@@ -35,6 +41,7 @@ func TestFieldSchema(t *testing.T) {
 		assert.Equal(t, field.AutoID, fieldSchema.GetAutoID())
 		assert.Equal(t, field.PrimaryKey, fieldSchema.GetIsPrimaryKey())
 		assert.Equal(t, field.IsPartitionKey, fieldSchema.GetIsPartitionKey())
+		assert.Equal(t, field.DefaultValue, fieldSchema.GetDefaultValue())
 		assert.Equal(t, field.IsDynamic, fieldSchema.GetIsDynamic())
 		assert.Equal(t, field.Description, fieldSchema.GetDescription())
 		assert.Equal(t, field.TypeParams, KvPairsMap(fieldSchema.GetTypeParams()))
@@ -49,6 +56,7 @@ func TestFieldSchema(t *testing.T) {
 		assert.Equal(t, field.Description, nf.Description)
 		assert.Equal(t, field.IsDynamic, nf.IsDynamic)
 		assert.Equal(t, field.IsPartitionKey, nf.IsPartitionKey)
+		assert.Equal(t, field.DefaultValue, nf.DefaultValue)
 		assert.EqualValues(t, field.TypeParams, nf.TypeParams)
 	}