From 6f8f2ee6a7a93256f217320e75b67dcff182beed Mon Sep 17 00:00:00 2001
From: Congqi Xia <congqi.xia@zilliz.com>
Date: Thu, 25 Jul 2024 17:19:13 +0800
Subject: [PATCH] enhance: Support specifying field schema to be cluster key

See also milvus-io/milvus#32939

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 entity/schema.go      | 52 +++++++++++++++++++++++++------------------
 entity/schema_test.go |  3 +++
 2 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/entity/schema.go b/entity/schema.go
index 03086bfa..9bb4754b 100644
--- a/entity/schema.go
+++ b/entity/schema.go
@@ -141,33 +141,35 @@ func (s *Schema) PKField() *Field {
 
 // Field represent field schema in milvus
 type Field struct {
-	ID             int64  // field id, generated when collection is created, input value is ignored
-	Name           string // field name
-	PrimaryKey     bool   // is primary key
-	AutoID         bool   // is auto id
-	Description    string
-	DataType       FieldType
-	TypeParams     map[string]string
-	IndexParams    map[string]string
-	IsDynamic      bool
-	IsPartitionKey bool
-	ElementType    FieldType
+	ID              int64  // field id, generated when collection is created, input value is ignored
+	Name            string // field name
+	PrimaryKey      bool   // is primary key
+	AutoID          bool   // is auto id
+	Description     string
+	DataType        FieldType
+	TypeParams      map[string]string
+	IndexParams     map[string]string
+	IsDynamic       bool
+	IsPartitionKey  bool
+	IsClusteringKey bool
+	ElementType     FieldType
 }
 
 // ProtoMessage generates corresponding FieldSchema
 func (f *Field) ProtoMessage() *schema.FieldSchema {
 	return &schema.FieldSchema{
-		FieldID:        f.ID,
-		Name:           f.Name,
-		Description:    f.Description,
-		IsPrimaryKey:   f.PrimaryKey,
-		AutoID:         f.AutoID,
-		DataType:       schema.DataType(f.DataType),
-		TypeParams:     MapKvPairs(f.TypeParams),
-		IndexParams:    MapKvPairs(f.IndexParams),
-		IsDynamic:      f.IsDynamic,
-		IsPartitionKey: f.IsPartitionKey,
-		ElementType:    schema.DataType(f.ElementType),
+		FieldID:         f.ID,
+		Name:            f.Name,
+		Description:     f.Description,
+		IsPrimaryKey:    f.PrimaryKey,
+		AutoID:          f.AutoID,
+		DataType:        schema.DataType(f.DataType),
+		TypeParams:      MapKvPairs(f.TypeParams),
+		IndexParams:     MapKvPairs(f.IndexParams),
+		IsDynamic:       f.IsDynamic,
+		IsPartitionKey:  f.IsPartitionKey,
+		IsClusteringKey: f.IsClusteringKey,
+		ElementType:     schema.DataType(f.ElementType),
 	}
 }
 
@@ -214,6 +216,11 @@ func (f *Field) WithIsPartitionKey(isPartitionKey bool) *Field {
 	return f
 }
 
+func (f *Field) WithIsClusteringKey(isClusteringKey bool) *Field {
+	f.IsClusteringKey = isClusteringKey
+	return f
+}
+
 /*
 func (f *Field) WithDefaultValueBool(defaultValue bool) *Field {
 	f.DefaultValue = &schema.ValueField{
@@ -318,6 +325,7 @@ func (f *Field) ReadProto(p *schema.FieldSchema) *Field {
 	f.IndexParams = KvPairsMap(p.GetIndexParams())
 	f.IsDynamic = p.GetIsDynamic()
 	f.IsPartitionKey = p.GetIsPartitionKey()
+	f.IsClusteringKey = p.GetIsClusteringKey()
 	f.ElementType = FieldType(p.GetElementType())
 
 	return f
diff --git a/entity/schema_test.go b/entity/schema_test.go
index 811a46ba..3a676436 100644
--- a/entity/schema_test.go
+++ b/entity/schema_test.go
@@ -28,6 +28,7 @@ func TestFieldSchema(t *testing.T) {
 		NewField().WithName("array_field").WithDataType(FieldTypeArray).WithElementType(FieldTypeBool).WithMaxCapacity(128),
 		NewField().WithName("fp16_field").WithDataType(FieldTypeFloat16Vector).WithDim(128),
 		NewField().WithName("bf16_field").WithDataType(FieldTypeBFloat16Vector).WithDim(128),
+		NewField().WithName("cluster_key").WithDataType(FieldTypeInt32).WithIsClusteringKey(true),
 		/*
 			NewField().WithName("default_value_bool").WithDataType(FieldTypeBool).WithDefaultValueBool(true),
 			NewField().WithName("default_value_int").WithDataType(FieldTypeInt32).WithDefaultValueInt(1),
@@ -45,6 +46,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.IsClusteringKey, fieldSchema.GetIsClusteringKey())
 		assert.Equal(t, field.IsDynamic, fieldSchema.GetIsDynamic())
 		assert.Equal(t, field.Description, fieldSchema.GetDescription())
 		assert.Equal(t, field.TypeParams, KvPairsMap(fieldSchema.GetTypeParams()))
@@ -60,6 +62,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.IsClusteringKey, nf.IsClusteringKey)
 		assert.EqualValues(t, field.TypeParams, nf.TypeParams)
 		assert.EqualValues(t, field.ElementType, nf.ElementType)
 	}