feat: Support Int8Vector in go

Signed-off-by: Cai Yudong <[email protected]>
milvus-io · Jan 6, 2025 · 41fcd1b · 41fcd1b
1 parent 731e882
commit 41fcd1b
Show file tree

Hide file tree

Showing 63 changed files with 1,600 additions and 203 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,7 @@ func FieldDataColumn(fd *schemapb.FieldData, begin, end int) (Column, error) {
 			vectors = append(vectors, vector)
 		}
 		return NewColumnSparseVectors(fd.GetFieldName(), vectors), nil
+
 	default:
 		return nil, fmt.Errorf("unsupported data type %s", fd.GetType())
 	}

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"
 	}
@@ -136,6 +140,8 @@ func (t FieldType) PbFieldType() (string, string) {
 		return "[]byte", ""
 	case FieldTypeBFloat16Vector:
 		return "[]byte", ""
+	case FieldTypeInt8Vector:
+		return "[]int8", ""
 	default:
 		return "undefined", ""
 	}
@@ -177,6 +183,8 @@ const (
 	FieldTypeBFloat16Vector FieldType = 103
 	// FieldTypeBinaryVector field type sparse vector
 	FieldTypeSparseVector FieldType = 104
+	// FieldTypeInt8Vector field type int8 vector
+	FieldTypeInt8Vector FieldType = 105
 )
 
 // Field represent field schema in milvus

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.

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-proto/go-api/v2 v2.5.0-beta.0.20250102080446-c3ba3d26a90f
 	github.com/milvus-io/milvus/pkg v0.0.2-0.20241126032235-cb6542339e84
 	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,8 +318,8 @@ 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-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.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/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=

diff --git a/go.mod b/go.mod
@@ -23,7 +23,7 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
 	github.com/klauspost/compress v1.17.9
 	github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d
-	github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20241211060635-410431d7865b
+	github.com/milvus-io/milvus-proto/go-api/v2 v2.5.0-beta.0.20250102080446-c3ba3d26a90f
 	github.com/minio/minio-go/v7 v7.0.73
 	github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81
 	github.com/prometheus/client_golang v1.14.0

diff --git a/go.sum b/go.sum
@@ -630,8 +630,8 @@ github.com/milvus-io/cgosymbolizer v0.0.0-20240722103217-b7dee0e50119 h1:9VXijWu
 github.com/milvus-io/cgosymbolizer v0.0.0-20240722103217-b7dee0e50119/go.mod h1:DvXTE/K/RtHehxU8/GtDs4vFtfw64jJ3PaCnFri8CRg=
 github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b h1:TfeY0NxYxZzUfIfYe5qYDBzt4ZYRqzUjTR6CvUzjat8=
 github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b/go.mod h1:iwW+9cWfIzzDseEBCCeDSN5SD16Tidvy8cwQ7ZY8Qj4=
-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-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/pulsar-client-go v0.12.1 h1:O2JZp1tsYiO7C0MQ4hrUY/aJXnn2Gry6hpm7UodghmE=
 github.com/milvus-io/pulsar-client-go v0.12.1/go.mod h1:dkutuH4oS2pXiGm+Ti7fQZ4MRjrMPZ8IJeEGAWMeckk=
 github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=

diff --git a/internal/core/src/bitset/bitset.h b/internal/core/src/bitset/bitset.h
@@ -989,11 +989,11 @@ class BitsetView : public BitsetBase<PolicyT,
 
     BitsetView() = default;
     BitsetView(const BitsetView&) = default;
-    BitsetView(BitsetView&&) noexcept = default;
+    BitsetView(BitsetView&&) = default;
     BitsetView&
     operator=(const BitsetView&) = default;
     BitsetView&
-    operator=(BitsetView&&) noexcept = default;
+    operator=(BitsetView&&) = default;
 
     template <typename ImplT, bool R>
     explicit BitsetView(BitsetBase<PolicyT, ImplT, R>& bitset)
@@ -1077,13 +1077,13 @@ class Bitset
     // Do not allow implicit copies (Rust style).
     Bitset(const Bitset&) = delete;
     // Allow default move.
-    Bitset(Bitset&&) noexcept = default;
+    Bitset(Bitset&&) = default;
     // Do not allow implicit copies (Rust style).
     Bitset&
     operator=(const Bitset&) = delete;
     // Allow default move.
     Bitset&
-    operator=(Bitset&&) noexcept = default;
+    operator=(Bitset&&) = default;
 
     template <typename C, bool R>
     explicit Bitset(const BitsetBase<PolicyT, C, R>& other) {

diff --git a/internal/core/src/common/Types.h b/internal/core/src/common/Types.h
@@ -58,6 +58,7 @@ using distance_t = float;
 using float16 = knowhere::fp16;
 using bfloat16 = knowhere::bf16;
 using bin1 = knowhere::bin1;
+using int8 = knowhere::int8;
 
 // See also: https://github.com/milvus-io/milvus-proto/blob/master/proto/schema.proto
 enum class DataType {
@@ -85,6 +86,7 @@ enum class DataType {
     VECTOR_FLOAT16 = 102,
     VECTOR_BFLOAT16 = 103,
     VECTOR_SPARSE_FLOAT = 104,
+    VECTOR_INT8 = 105,
 };
 
 using Timestamp = uint64_t;  // TODO: use TiKV-like timestamp
@@ -322,6 +324,11 @@ IsSparseFloatVectorDataType(DataType data_type) {
     return data_type == DataType::VECTOR_SPARSE_FLOAT;
 }
 
+inline bool
+IsInt8VectorDataType(DataType data_type) {
+    return data_type == DataType::VECTOR_INT8;
+}
+
 inline bool
 IsFloatVectorDataType(DataType data_type) {
     return IsDenseFloatVectorDataType(data_type) ||
@@ -331,7 +338,7 @@ IsFloatVectorDataType(DataType data_type) {
 inline bool
 IsVectorDataType(DataType data_type) {
     return IsBinaryVectorDataType(data_type) ||
-           IsFloatVectorDataType(data_type);
+           IsFloatVectorDataType(data_type) || IsInt8VectorDataType(data_type);
 }
 
 inline bool
@@ -418,7 +425,17 @@ IsFloatVectorMetricType(const MetricType& metric_type) {
 
 inline bool
 IsBinaryVectorMetricType(const MetricType& metric_type) {
-    return !IsFloatVectorMetricType(metric_type);
+    return metric_type == knowhere::metric::HAMMING ||
+           metric_type == knowhere::metric::JACCARD ||
+           metric_type == knowhere::metric::SUPERSTRUCTURE ||
+           metric_type == knowhere::metric::SUBSTRUCTURE;
+}
+
+inline bool
+IsIntVectorMetricType(const MetricType& metric_type) {
+    return metric_type == knowhere::metric::L2 ||
+           metric_type == knowhere::metric::IP ||
+           metric_type == knowhere::metric::COSINE;
 }
 
 // Plus 1 because we can't use greater(>) symbol

diff --git a/internal/core/src/index/IndexFactory.cpp b/internal/core/src/index/IndexFactory.cpp
@@ -443,6 +443,10 @@ IndexFactory::CreateVectorIndex(
                 return std::make_unique<VectorMemIndex<bfloat16>>(
                     index_type, metric_type, version, file_manager_context);
             }
+            case DataType::VECTOR_INT8: {
+                return std::make_unique<VectorMemIndex<int8>>(
+                    index_type, metric_type, version, file_manager_context);
+            }
             default:
                 PanicInfo(
                     DataTypeInvalid,

diff --git a/internal/core/src/index/Utils.cpp b/internal/core/src/index/Utils.cpp
@@ -63,6 +63,7 @@ BIN_List() {
     return ret;
 }
 
+// TODO caiyd: should list supported list
 std::vector<std::tuple<IndexType, MetricType>>
 unsupported_index_combinations() {
     static std::vector<std::tuple<IndexType, MetricType>> ret{

diff --git a/internal/core/src/index/Utils.h b/internal/core/src/index/Utils.h
@@ -105,11 +105,13 @@ CheckMetricTypeSupport(const MetricType& metric_type) {
     if constexpr (std::is_same_v<T, bin1>) {
         AssertInfo(
             IsBinaryVectorMetricType(metric_type),
-            "binary vector does not float vector metric type: " + metric_type);
+            "binary vector does not support metric type: " + metric_type);
+    } else if constexpr (std::is_same_v<T, int8>) {
+        AssertInfo(IsIntVectorMetricType(metric_type),
+                   "int vector does not support metric type: " + metric_type);
     } else {
-        AssertInfo(
-            IsFloatVectorMetricType(metric_type),
-            "float vector does not binary vector metric type: " + metric_type);
+        AssertInfo(IsFloatVectorMetricType(metric_type),
+                   "float vector does not support metric type: " + metric_type);
     }
 }
 

diff --git a/internal/core/src/index/VectorMemIndex.cpp b/internal/core/src/index/VectorMemIndex.cpp
@@ -636,5 +636,6 @@ template class VectorMemIndex<float>;
 template class VectorMemIndex<bin1>;
 template class VectorMemIndex<float16>;
 template class VectorMemIndex<bfloat16>;
+template class VectorMemIndex<int8>;
 
 }  // namespace milvus::index
diff --git a/internal/core/src/indexbuilder/IndexFactory.h b/internal/core/src/indexbuilder/IndexFactory.h
@@ -67,6 +67,7 @@ class IndexFactory {
             case DataType::VECTOR_BFLOAT16:
             case DataType::VECTOR_BINARY:
             case DataType::VECTOR_SPARSE_FLOAT:
+            case DataType::VECTOR_INT8:
                 return std::make_unique<VecIndexCreator>(type, config, context);
             default:
                 PanicInfo(DataTypeInvalid,

diff --git a/internal/core/src/indexbuilder/index_c.cpp b/internal/core/src/indexbuilder/index_c.cpp
@@ -450,6 +450,29 @@ BuildSparseFloatVecIndex(CIndex index,
     return status;
 }
 
+CStatus
+BuildInt8VecIndex(CIndex index, int64_t int8_value_num, const int8_t* vectors) {
+    auto status = CStatus();
+    try {
+        AssertInfo(index,
+                   "failed to build int8 vector index, passed index was null");
+        auto real_index =
+            reinterpret_cast<milvus::indexbuilder::IndexCreatorBase*>(index);
+        auto cIndex =
+            dynamic_cast<milvus::indexbuilder::VecIndexCreator*>(real_index);
+        auto dim = cIndex->dim();
+        auto row_nums = int8_value_num / dim;
+        auto ds = knowhere::GenDataSet(row_nums, dim, vectors);
+        cIndex->Build(ds);
+        status.error_code = Success;
+        status.error_msg = "";
+    } catch (std::exception& e) {
+        status.error_code = UnexpectedError;
+        status.error_msg = strdup(e.what());
+    }
+    return status;
+}
+
 // field_data:
 //  1, serialized proto::schema::BoolArray, if type is bool;
 //  2, serialized proto::schema::StringArray, if type is string;

diff --git a/internal/core/src/indexbuilder/index_c.h b/internal/core/src/indexbuilder/index_c.h
@@ -58,6 +58,9 @@ BuildSparseFloatVecIndex(CIndex index,
                          int64_t dim,
                          const uint8_t* vectors);
 
+CStatus
+BuildInt8VecIndex(CIndex index, int64_t data_size, const int8_t* vectors);
+
 // field_data:
 //  1, serialized proto::schema::BoolArray, if type is bool;
 //  2, serialized proto::schema::StringArray, if type is string;

diff --git a/internal/core/src/segcore/vector_index_c.cpp b/internal/core/src/segcore/vector_index_c.cpp
@@ -73,6 +73,12 @@ ValidateIndexParams(const char* index_type,
                 knowhere::Version::GetCurrentVersion().VersionNumber(),
                 json,
                 error_msg);
+        } else if (dataType == milvus::DataType::VECTOR_INT8) {
+            status = knowhere::IndexStaticFaced<knowhere::int8>::ConfigCheck(
+                index_type,
+                knowhere::Version::GetCurrentVersion().VersionNumber(),
+                json,
+                error_msg);
         } else {
             status = knowhere::Status::invalid_args;
         }

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]
 						}
 					]
 				}`),