feat: Support restful & go sdk for Int8Vector

Signed-off-by: Cai Yudong <[email protected]>
milvus-io · Jan 15, 2025 · 3747fd7 · 3747fd7
1 parent 3a6408b
commit 3747fd7
Show file tree

Hide file tree

Showing 16 changed files with 413 additions and 32 deletions.
diff --git a/client/column/columns.go b/client/column/columns.go
@@ -284,6 +284,7 @@ func FieldDataColumn(fd *schemapb.FieldData, begin, end int) (Column, error) {
 			vector = append(vector, v)
 		}
 		return NewColumnBFloat16Vector(fd.GetFieldName(), dim, vector), nil
+
 	case schemapb.DataType_SparseFloatVector:
 		sparseVectors := fd.GetVectors().GetSparseFloatVector()
 		if sparseVectors == nil {
@@ -303,6 +304,29 @@ func FieldDataColumn(fd *schemapb.FieldData, begin, end int) (Column, error) {
 			vectors = append(vectors, vector)
 		}
 		return NewColumnSparseVectors(fd.GetFieldName(), vectors), nil
+
+	case schemapb.DataType_Int8Vector:
+		vectors := fd.GetVectors()
+		x, ok := vectors.GetData().(*schemapb.VectorField_Int8Vector)
+		if !ok {
+			return nil, errFieldDataTypeNotMatch
+		}
+		data := x.Int8Vector
+		dim := int(vectors.GetDim())
+		if end < 0 {
+			end = len(data) / dim
+		}
+		vector := make([][]int8, 0, end-begin) // shall not have remanunt
+		// TODO caiyd: has better way to convert []byte to []int8 ?
+		for i := begin; i < end; i++ {
+			v := make([]int8, dim)
+			for j := 0; j < dim; j++ {
+				v[j] = int8(data[i*dim+j])
+			}
+			vector = append(vector, v)
+		}
+		return NewColumnInt8Vector(fd.GetFieldName(), dim, vector), nil
+
 	default:
 		return nil, fmt.Errorf("unsupported data type %s", fd.GetType())
 	}

diff --git a/client/column/vector.go b/client/column/vector.go
@@ -213,3 +213,36 @@ func (c *ColumnBFloat16Vector) Slice(start, end int) Column {
 		vectorBase: c.vectorBase.slice(start, end),
 	}
 }
+
+/* int8 vector */
+
+type ColumnInt8Vector struct {
+	*vectorBase[entity.Int8Vector]
+}
+
+func NewColumnInt8Vector(fieldName string, dim int, data [][]int8) *ColumnInt8Vector {
+	vectors := lo.Map(data, func(row []int8, _ int) entity.Int8Vector { return entity.Int8Vector(row) })
+	return &ColumnInt8Vector{
+		vectorBase: newVectorBase(fieldName, dim, vectors, entity.FieldTypeInt8Vector),
+	}
+}
+
+// AppendValue appends vector value into values.
+// override default type constrains, add `[]int8` conversion
+func (c *ColumnInt8Vector) AppendValue(i interface{}) error {
+	switch vector := i.(type) {
+	case entity.Int8Vector:
+		c.values = append(c.values, vector)
+	case []int8:
+		c.values = append(c.values, vector)
+	default:
+		return errors.Newf("unexpected append value type %T, field type %v", vector, c.fieldType)
+	}
+	return nil
+}
+
+func (c *ColumnInt8Vector) Slice(start, end int) Column {
+	return &ColumnInt8Vector{
+		vectorBase: c.vectorBase.slice(start, end),
+	}
+}
diff --git a/client/column/vector_test.go b/client/column/vector_test.go
@@ -187,6 +187,38 @@ func (s *VectorSuite) TestBasic() {
 			}
 		}
 	})
+
+	s.Run("int8_vector", func() {
+		name := fmt.Sprintf("field_%d", rand.Intn(1000))
+		n := 3
+		dim := rand.Intn(10) + 2
+		data := make([][]int8, 0, n)
+		for i := 0; i < n; i++ {
+			row := lo.RepeatBy(dim, func(i int) int8 {
+				return int8(rand.Intn(256) - 128)
+			})
+			data = append(data, row)
+		}
+		column := NewColumnInt8Vector(name, dim, data)
+		s.Equal(entity.FieldTypeInt8Vector, column.Type())
+		s.Equal(name, column.Name())
+		s.Equal(lo.Map(data, func(row []int8, _ int) entity.Int8Vector { return entity.Int8Vector(row) }), column.Data())
+		s.Equal(dim, column.Dim())
+
+		fd := column.FieldData()
+		s.Equal(name, fd.GetFieldName())
+		s.Equal(lo.Flatten(data), fd.GetVectors().GetInt8Vector())
+
+		result, err := FieldDataColumn(fd, 0, -1)
+		s.NoError(err)
+		parsed, ok := result.(*ColumnInt8Vector)
+		if s.True(ok) {
+			s.Equal(entity.FieldTypeInt8Vector, parsed.Type())
+			s.Equal(name, parsed.Name())
+			s.Equal(lo.Map(data, func(row []int8, _ int) entity.Int8Vector { return entity.Int8Vector(row) }), parsed.Data())
+			s.Equal(dim, parsed.Dim())
+		}
+	})
 }
 
 func (s *VectorSuite) TestSlice() {
@@ -277,6 +309,28 @@ func (s *VectorSuite) TestSlice() {
 			s.Equal(lo.Map(data[:l], func(row []byte, _ int) entity.BFloat16Vector { return entity.BFloat16Vector(row) }), slicedColumn.Data())
 		}
 	})
+
+	s.Run("int8_vector", func() {
+		name := fmt.Sprintf("field_%d", rand.Intn(1000))
+		n := 100
+		dim := rand.Intn(10) + 2
+		data := make([][]int8, 0, n)
+		for i := 0; i < n; i++ {
+			row := lo.RepeatBy(dim, func(i int) int8 {
+				return int8(rand.Intn(256) - 128)
+			})
+			data = append(data, row)
+		}
+		column := NewColumnInt8Vector(name, dim, data)
+
+		l := rand.Intn(n)
+		sliced := column.Slice(0, l)
+		slicedColumn, ok := sliced.(*ColumnInt8Vector)
+		if s.True(ok) {
+			s.Equal(dim, slicedColumn.Dim())
+			s.Equal(lo.Map(data[:l], func(row []int8, _ int) entity.Int8Vector { return entity.Int8Vector(row) }), slicedColumn.Data())
+		}
+	})
 }
 
 func TestVectors(t *testing.T) {

diff --git a/client/entity/field.go b/client/entity/field.go
@@ -62,6 +62,8 @@ func (t FieldType) Name() string {
 		return "Float16Vector"
 	case FieldTypeBFloat16Vector:
 		return "BFloat16Vector"
+	case FieldTypeInt8Vector:
+		return "Int8Vector"
 	default:
 		return "undefined"
 	}
@@ -100,6 +102,8 @@ func (t FieldType) String() string {
 		return "[]byte"
 	case FieldTypeBFloat16Vector:
 		return "[]byte"
+	case FieldTypeInt8Vector:
+		return "[]int8"
 	default:
 		return "undefined"
 	}

diff --git a/client/entity/vectors.go b/client/entity/vectors.go
@@ -56,7 +56,7 @@ func (fv FloatVector) ToBFloat16Vector() BFloat16Vector {
 	return typeutil.Float32ArrayToBFloat16Bytes(fv)
 }
 
-// FloatVector float32 vector wrapper.
+// Float16Vector float16 vector wrapper.
 type Float16Vector []byte
 
 // Dim returns vector dimension.
@@ -77,7 +77,7 @@ func (fv Float16Vector) ToFloat32Vector() FloatVector {
 	return typeutil.Float16BytesToFloat32Vector(fv)
 }
 
-// FloatVector float32 vector wrapper.
+// BFloat16Vector bfloat16 vector wrapper.
 type BFloat16Vector []byte
 
 // Dim returns vector dimension.
@@ -131,3 +131,21 @@ func (t Text) FieldType() FieldType {
 func (t Text) Serialize() []byte {
 	return []byte(t)
 }
+
+// Int8Vector []int8 vector wrapper
+type Int8Vector []int8
+
+// Dim return vector dimension
+func (iv Int8Vector) Dim() int {
+	return len(iv)
+}
+
+// Serialize just return bytes
+func (iv Int8Vector) Serialize() []byte {
+	return typeutil.Int8ArrayToBytes(iv)
+}
+
+// entity.FieldType returns coresponding field type.
+func (iv Int8Vector) FieldType() FieldType {
+	return FieldTypeInt8Vector
+}
diff --git a/client/entity/vectors_test.go b/client/entity/vectors_test.go
@@ -92,4 +92,15 @@ func TestVectors(t *testing.T) {
 		assert.Equal(t, dim*8, bv.Dim())
 		assert.ElementsMatch(t, raw, bv.Serialize())
 	})
+
+	t.Run("test int8 vector", func(t *testing.T) {
+		raw := make([]int8, dim)
+		for i := 0; i < dim; i++ {
+			raw[i] = int8(rand.Intn(256) - 128)
+		}
+
+		iv := Int8Vector(raw)
+		assert.Equal(t, dim, iv.Dim())
+		assert.Equal(t, dim, len(iv.Serialize()))
+	})
 }
diff --git a/client/go.mod b/client/go.mod
@@ -6,14 +6,13 @@ require (
 	github.com/blang/semver/v4 v4.0.0
 	github.com/cockroachdb/errors v1.9.1
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
-	github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20241211060635-410431d7865b
-	github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84
+	github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20250102080446-c3ba3d26a90f
+	github.com/milvus-io/milvus/pkg v0.0.2-0.20250115044500-f5234c3c11a3
 	github.com/quasilyte/go-ruleguard/dsl v0.3.22
 	github.com/samber/lo v1.27.0
 	github.com/stretchr/testify v1.9.0
 	github.com/tidwall/gjson v1.17.1
 	go.uber.org/atomic v1.10.0
-	golang.org/x/exp v0.0.0-20230224173230-c95f2b4c22f2
 	google.golang.org/grpc v1.65.0
 	google.golang.org/protobuf v1.34.2
 )
@@ -99,6 +98,7 @@ require (
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.27.0 // indirect
 	golang.org/x/crypto v0.31.0 // indirect
+	golang.org/x/exp v0.0.0-20230224173230-c95f2b4c22f2 // indirect
 	golang.org/x/net v0.33.0 // indirect
 	golang.org/x/sync v0.10.0 // indirect
 	golang.org/x/sys v0.28.0 // indirect

diff --git a/client/go.sum b/client/go.sum
@@ -318,10 +318,10 @@ github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfr
 github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8=
 github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc=
 github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg=
-github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20241211060635-410431d7865b h1:iPPhnFx+s7FF53UeWj7A4EYhPRMFPL6mHqyQw7qRjeQ=
-github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20241211060635-410431d7865b/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
-github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84 h1:EAFxmxUVp5yYFDCrX1MQoSxkTO+ycy8NXEqEDEB3cRM=
-github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84/go.mod h1:RATa0GS4jhkPpsYOvQ/QvcNz8rd+TlRPDiSyXQnMMxs=
+github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20250102080446-c3ba3d26a90f h1:So6RKU5wqP/8EaKogicJP8gZ2SrzzS/JprusBaE3RKc=
+github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20250102080446-c3ba3d26a90f/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
+github.com/milvus-io/milvus/pkg v0.0.2-0.20250115044500-f5234c3c11a3 h1:WF9BkNk1XjLtwMbaB/cniRBMMNLnqG6e+AUbK8DciHQ=
+github.com/milvus-io/milvus/pkg v0.0.2-0.20250115044500-f5234c3c11a3/go.mod h1:nxnHkDFB3jh27nTQJBaC4azAQO8chT03DkmoiZ5086s=
 github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=
 github.com/mitchellh/go-homedir v1.0.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=

diff --git a/internal/distributed/proxy/httpserver/handler_v1_test.go b/internal/distributed/proxy/httpserver/handler_v1_test.go
@@ -1270,7 +1270,8 @@ func TestFp16Bf16VectorsV1(t *testing.T) {
 							"binaryVector": "AQ==",
 							"float16Vector": [3.0],
 							"bfloat16Vector": [4.4, 442],
-							"sparseFloatVector": {"1": 0.1, "2": 0.44}
+							"sparseFloatVector": {"1": 0.1, "2": 0.44},
+							"int8Vector": [1, 2]
 						}
 					]
 				}`),
@@ -1289,7 +1290,8 @@ func TestFp16Bf16VectorsV1(t *testing.T) {
 							"binaryVector": "AQ==",
 							"float16Vector": [3, 3.0],
 							"bfloat16Vector": [4.4, 442],
-							"sparseFloatVector": {"1": 0.1, "2": 0.44}
+							"sparseFloatVector": {"1": 0.1, "2": 0.44},
+							"int8Vector": [1, 2]
 						}
 					]
 				}`),
@@ -1308,7 +1310,8 @@ func TestFp16Bf16VectorsV1(t *testing.T) {
 							"binaryVector": "AQ==",
 							"float16Vector": [3, 3],
 							"bfloat16Vector": [4.4, 442],
-							"sparseFloatVector": {"1": 0.1, "2": 0.44}
+							"sparseFloatVector": {"1": 0.1, "2": 0.44},
+							"int8Vector": [1, 2]
 						}
 					]
 				}`),
@@ -1326,7 +1329,8 @@ func TestFp16Bf16VectorsV1(t *testing.T) {
 							"binaryVector": "AQ==",
 							"float16Vector": "AQIDBA==",
 							"bfloat16Vector": "AQIDBA==",
-							"sparseFloatVector": {"1": 0.1, "2": 0.44}
+							"sparseFloatVector": {"1": 0.1, "2": 0.44},
+							"int8Vector": [1, 2]
 						}
 					]
 				}`),
@@ -1344,7 +1348,8 @@ func TestFp16Bf16VectorsV1(t *testing.T) {
 							"binaryVector": "AQ==",
 							"float16Vector": [3, 3.0, 3],
 							"bfloat16Vector": [4.4, 44],
-							"sparseFloatVector": {"1": 0.1, "2": 0.44}
+							"sparseFloatVector": {"1": 0.1, "2": 0.44},
+							"int8Vector": [1, 2]
 						}
 					]
 				}`),
@@ -1363,7 +1368,8 @@ func TestFp16Bf16VectorsV1(t *testing.T) {
 							"binaryVector": "AQ==",
 							"float16Vector": [3, 3.0],
 							"bfloat16Vector": [4.4, 442, 44],
-							"sparseFloatVector": {"1": 0.1, "2": 0.44}
+							"sparseFloatVector": {"1": 0.1, "2": 0.44},
+							"int8Vector": [1, 2]
 						}
 					]
 				}`),
@@ -1382,7 +1388,8 @@ func TestFp16Bf16VectorsV1(t *testing.T) {
 							"binaryVector": "AQ==",
 							"float16Vector": "AQIDBA==",
 							"bfloat16Vector": [4.4, 442],
-							"sparseFloatVector": {"1": 0.1, "2": 0.44}
+							"sparseFloatVector": {"1": 0.1, "2": 0.44},
+							"int8Vector": [1, 2]
 						},
 						{
 							"book_id": 1,
@@ -1391,7 +1398,8 @@ func TestFp16Bf16VectorsV1(t *testing.T) {
 							"binaryVector": "AQ==",
 							"float16Vector": [3.1, 3.1],
 							"bfloat16Vector": "AQIDBA==",
-							"sparseFloatVector": {"3": 1.1, "2": 0.44}
+							"sparseFloatVector": {"3": 1.1, "2": 0.44},
+							"int8Vector": [1, 2]
 						}
 					]
 				}`),