diff --git a/internal/core/src/segcore/Utils.cpp b/internal/core/src/segcore/Utils.cpp index c33fe01909..ba55ca8706 100644 --- a/internal/core/src/segcore/Utils.cpp +++ b/internal/core/src/segcore/Utils.cpp @@ -359,7 +359,7 @@ CreateVectorDataArray(int64_t count, const FieldMeta& field_meta) { case DataType::VECTOR_INT8: { auto length = count * dim; auto obj = vector_array->mutable_int8_vector(); - obj->resize(length * sizeof(int8)); + obj->resize(length); break; } default: { diff --git a/internal/core/unittest/test_chunk_vector.cpp b/internal/core/unittest/test_chunk_vector.cpp index 66867bbbd0..85ccc41ab9 100644 --- a/internal/core/unittest/test_chunk_vector.cpp +++ b/internal/core/unittest/test_chunk_vector.cpp @@ -10,6 +10,7 @@ // or implied. See the License for the specific language governing permissions and limitations under the License #include +#include #include "common/Types.h" #include "knowhere/comp/index_param.h" @@ -71,6 +72,8 @@ TEST_F(ChunkVectorTest, FillDataWithMmap) { "bf16_vec", DataType::VECTOR_BFLOAT16, 128, metric_type); auto sparse_vec = schema->AddDebugField( "sparse_vec", DataType::VECTOR_SPARSE_FLOAT, 128, metric_type); + auto int8_vec = schema->AddDebugField( + "int8_vec", DataType::VECTOR_INT8, 128, metric_type); schema->set_primary_field_id(int64_field); std::map index_params = { @@ -136,6 +139,8 @@ TEST_F(ChunkVectorTest, FillDataWithMmap) { segment->bulk_subscript(bf16_vec, ids_ds->GetIds(), num_inserted); auto sparse_vec_result = segment->bulk_subscript(sparse_vec, ids_ds->GetIds(), num_inserted); + auto int8_vec_result = + segment->bulk_subscript(int8_vec, ids_ds->GetIds(), num_inserted); EXPECT_EQ(bool_result->scalars().bool_data().data_size(), num_inserted); EXPECT_EQ(int8_result->scalars().int_data().data_size(), num_inserted); @@ -159,6 +164,8 @@ TEST_F(ChunkVectorTest, FillDataWithMmap) { EXPECT_EQ( sparse_vec_result->vectors().sparse_float_vector().contents_size(), num_inserted); + EXPECT_EQ(int8_vec_result->vectors().int8_vector().size(), + num_inserted * dim); EXPECT_EQ(int_array_result->scalars().array_data().data_size(), num_inserted); EXPECT_EQ(long_array_result->scalars().array_data().data_size(), @@ -184,24 +191,33 @@ TEST_F(ChunkVectorTest, FillDataWithMmap) { .data(); auto sparse_vec_res = SparseBytesToRows( sparse_vec_result->vectors().sparse_float_vector().contents()); + auto int8_vec_res = (int8*)int8_vec_result.get() + ->mutable_vectors() + ->int8_vector() + .data(); EXPECT_TRUE(fp32_vec_res.size() == num_inserted * dim); auto fp32_vec_gt = dataset.get_col(fp32_vec); auto fp16_vec_gt = dataset.get_col(fp16_vec); auto bf16_vec_gt = dataset.get_col(bf16_vec); auto sparse_vec_gt = dataset.get_col>(sparse_vec); + auto int8_vec_gt = dataset.get_col(int8_vec); for (size_t i = 0; i < num_inserted; ++i) { auto id = ids_ds->GetIds()[i]; // check dense vector - for (size_t j = 0; j < 128; ++j) { - EXPECT_TRUE(fp32_vec_res[i * dim + j] == - fp32_vec_gt[(id % per_batch) * dim + j]); - EXPECT_TRUE(fp16_vec_res[i * dim + j] == - fp16_vec_gt[(id % per_batch) * dim + j]); - EXPECT_TRUE(bf16_vec_res[i * dim + j] == - bf16_vec_gt[(id % per_batch) * dim + j]); - } + EXPECT_TRUE(memcmp((void*)(&fp32_vec_res[i * dim]), + (void*)(&fp32_vec_gt[(id % per_batch) * dim]), + sizeof(float) * dim) == 0); + EXPECT_TRUE(memcmp((void*)(&fp16_vec_res[i * dim]), + (void*)(&fp16_vec_gt[(id % per_batch) * dim]), + sizeof(float16) * dim) == 0); + EXPECT_TRUE(memcmp((void*)(&bf16_vec_res[i * dim]), + (void*)(&bf16_vec_gt[(id % per_batch) * dim]), + sizeof(bfloat16) * dim) == 0); + EXPECT_TRUE(memcmp((void*)(&int8_vec_res[i * dim]), + (void*)(&int8_vec_gt[(id % per_batch) * dim]), + sizeof(int8) * dim) == 0); //check sparse vector auto actual_row = sparse_vec_res[i]; auto expected_row = sparse_vec_gt[(id % per_batch)]; diff --git a/internal/core/unittest/test_loading.cpp b/internal/core/unittest/test_loading.cpp index 851663a3d5..92403e24b4 100644 --- a/internal/core/unittest/test_loading.cpp +++ b/internal/core/unittest/test_loading.cpp @@ -49,6 +49,8 @@ class IndexLoadTest : public ::testing::TestWithParam { data_type = milvus::DataType::VECTOR_BINARY; } else if (field_type == "vector_sparse_float") { data_type = milvus::DataType::VECTOR_SPARSE_FLOAT; + } else if (field_type == "vector_int8") { + data_type = milvus::DataType::VECTOR_INT8; } else if (field_type == "array") { data_type = milvus::DataType::ARRAY; } else { @@ -106,6 +108,22 @@ INSTANTIATE_TEST_SUITE_P( {"mmap", "true"}, {"field_type", "vector_fp16"}}, {0.125f, 1.0f, 0.0f, 1.0f, true}), + std::pair, LoadResourceRequest>( + {{"index_type", "HNSW"}, + {"metric_type", "L2"}, + {"efConstrcution", "300"}, + {"M", "30"}, + {"mmap", "false"}, + {"field_type", "vector_int8"}}, + {2.0f, 0.0f, 1.0f, 0.0f, true}), + std::pair, LoadResourceRequest>( + {{"index_type", "HNSW"}, + {"metric_type", "L2"}, + {"efConstrcution", "300"}, + {"M", "30"}, + {"mmap", "true"}, + {"field_type", "vector_int8"}}, + {0.125f, 1.0f, 0.0f, 1.0f, true}), std::pair, LoadResourceRequest>( {{"index_type", "IVFFLAT"}, {"metric_type", "L2"}, diff --git a/internal/core/unittest/test_sealed.cpp b/internal/core/unittest/test_sealed.cpp index 15e7f50b8a..e529736f48 100644 --- a/internal/core/unittest/test_sealed.cpp +++ b/internal/core/unittest/test_sealed.cpp @@ -2197,6 +2197,8 @@ TEST(Sealed, QueryAllFields) { "float16_vec", DataType::VECTOR_FLOAT16, 128, metric_type); auto bfloat16_vec = schema->AddDebugField( "bfloat16_vec", DataType::VECTOR_BFLOAT16, 128, metric_type); + auto int8_vec = schema->AddDebugField( + "int8_vec", DataType::VECTOR_INT8, 128, metric_type); schema->set_primary_field_id(int64_field); std::map index_params = { @@ -2235,6 +2237,7 @@ TEST(Sealed, QueryAllFields) { auto vector_values = dataset.get_col(vec); auto float16_vector_values = dataset.get_col(float16_vec); auto bfloat16_vector_values = dataset.get_col(bfloat16_vec); + auto int8_vector_values = dataset.get_col(int8_vec); auto ids_ds = GenRandomIds(dataset_size); auto bool_result = @@ -2273,6 +2276,8 @@ TEST(Sealed, QueryAllFields) { segment->bulk_subscript(float16_vec, ids_ds->GetIds(), dataset_size); auto bfloat16_vec_result = segment->bulk_subscript(bfloat16_vec, ids_ds->GetIds(), dataset_size); + auto int8_vec_result = + segment->bulk_subscript(int8_vec, ids_ds->GetIds(), dataset_size); EXPECT_EQ(bool_result->scalars().bool_data().data_size(), dataset_size); EXPECT_EQ(int8_result->scalars().int_data().data_size(), dataset_size); @@ -2290,6 +2295,8 @@ TEST(Sealed, QueryAllFields) { dataset_size * dim * 2); EXPECT_EQ(bfloat16_vec_result->vectors().bfloat16_vector().size(), dataset_size * dim * 2); + EXPECT_EQ(int8_vec_result->vectors().int8_vector().size(), + dataset_size * dim); EXPECT_EQ(int_array_result->scalars().array_data().data_size(), dataset_size); EXPECT_EQ(long_array_result->scalars().array_data().data_size(), diff --git a/internal/core/unittest/test_utils/DataGen.h b/internal/core/unittest/test_utils/DataGen.h index 24a82ffe94..f1004729f4 100644 --- a/internal/core/unittest/test_utils/DataGen.h +++ b/internal/core/unittest/test_utils/DataGen.h @@ -429,6 +429,7 @@ inline GeneratedData DataGen(SchemaPtr schema, case DataType::VECTOR_INT8: { auto dim = field_meta.get_dim(); vector final(dim * N); + srand(seed); for (auto& x : final) { x = int8_t(rand() % 256 - 128); } diff --git a/internal/parser/planparserv2/plan_parser_v2.go b/internal/parser/planparserv2/plan_parser_v2.go index 89ae5344e4..caacb7b6bb 100644 --- a/internal/parser/planparserv2/plan_parser_v2.go +++ b/internal/parser/planparserv2/plan_parser_v2.go @@ -191,6 +191,8 @@ func CreateSearchPlan(schema *typeutil.SchemaHelper, exprStr string, vectorField vectorType = planpb.VectorType_BFloat16Vector case schemapb.DataType_SparseFloatVector: vectorType = planpb.VectorType_SparseFloatVector + case schemapb.DataType_Int8Vector: + vectorType = planpb.VectorType_Int8Vector default: log.Error("Invalid dataType", zap.Any("dataType", dataType)) return nil, err diff --git a/internal/storage/serde.go b/internal/storage/serde.go index ae46ca8bfb..f0919646bf 100644 --- a/internal/storage/serde.go +++ b/internal/storage/serde.go @@ -459,6 +459,13 @@ var serdeMap = func() map[schemapb.DataType]serdeEntry { fixedSizeDeserializer, fixedSizeSerializer, } + m[schemapb.DataType_Int8Vector] = serdeEntry{ + func(i int) arrow.DataType { + return &arrow.FixedSizeBinaryType{ByteWidth: i} + }, + fixedSizeDeserializer, + fixedSizeSerializer, + } m[schemapb.DataType_FloatVector] = serdeEntry{ func(i int) arrow.DataType { return &arrow.FixedSizeBinaryType{ByteWidth: i * 4} diff --git a/tests/integration/getvector/get_vector_test.go b/tests/integration/getvector/get_vector_test.go index 7ee16d68aa..e729213a59 100644 --- a/tests/integration/getvector/get_vector_test.go +++ b/tests/integration/getvector/get_vector_test.go @@ -132,6 +132,8 @@ func (s *TestGetVectorSuite) run() { vecFieldData = integration.NewBFloat16VectorFieldData(vecFieldName, NB, dim) } else if typeutil.IsSparseFloatVectorType(s.vecType) { vecFieldData = integration.NewSparseFloatVectorFieldData(vecFieldName, NB) + } else if s.vecType == schemapb.DataType_Int8Vector { + vecFieldData = integration.NewInt8VectorFieldData(vecFieldName, NB, dim) } else { vecFieldData = integration.NewBinaryVectorFieldData(vecFieldName, NB, dim) } @@ -294,6 +296,26 @@ func (s *TestGetVectorSuite) run() { s.Require().Equal(rawData[id], resData[i]) } } + } else if s.vecType == schemapb.DataType_Int8Vector { + s.Require().Len(result.GetFieldsData()[vecFieldIndex].GetVectors().GetInt8Vector(), nq*topk*dim) + rawData := vecFieldData.GetVectors().GetInt8Vector() + resData := result.GetFieldsData()[vecFieldIndex].GetVectors().GetInt8Vector() + rowBytes := dim + if s.pkType == schemapb.DataType_Int64 { + for i, id := range result.GetIds().GetIntId().GetData() { + expect := rawData[int(id)*rowBytes : (int(id)+1)*rowBytes] + actual := resData[i*rowBytes : (i+1)*rowBytes] + s.Require().ElementsMatch(expect, actual) + } + } else { + for i, idStr := range result.GetIds().GetStrId().GetData() { + id, err := strconv.Atoi(idStr) + s.Require().NoError(err) + expect := rawData[id*rowBytes : (id+1)*rowBytes] + actual := resData[i*rowBytes : (i+1)*rowBytes] + s.Require().ElementsMatch(expect, actual) + } + } } else { s.Require().Len(result.GetFieldsData()[vecFieldIndex].GetVectors().GetBinaryVector(), nq*topk*dim/8) rawData := vecFieldData.GetVectors().GetBinaryVector() @@ -448,6 +470,16 @@ func (s *TestGetVectorSuite) TestGetVector_BFloat16Vector() { s.run() } +func (s *TestGetVectorSuite) TestGetVector_Int8Vector() { + s.nq = 10 + s.topK = 10 + s.indexType = integration.IndexHNSW + s.metricType = metric.L2 + s.pkType = schemapb.DataType_Int64 + s.vecType = schemapb.DataType_Int8Vector + s.run() +} + func (s *TestGetVectorSuite) TestGetVector_Big_NQ_TOPK() { s.T().Skip("skip big NQ Top due to timeout") s.nq = 10000 diff --git a/tests/integration/import/import_test.go b/tests/integration/import/import_test.go index 2cfe97f18a..a979cd1186 100644 --- a/tests/integration/import/import_test.go +++ b/tests/integration/import/import_test.go @@ -244,6 +244,11 @@ func (s *BulkInsertSuite) TestMultiFileTypes() { s.metricType = metric.L2 s.run() + s.vecType = schemapb.DataType_Int8Vector + s.indexType = "HNSW" + s.metricType = metric.L2 + s.run() + // TODO: not support numpy for SparseFloatVector by now if fileType != importutilv2.Numpy { s.vecType = schemapb.DataType_SparseFloatVector diff --git a/tests/integration/import/util_test.go b/tests/integration/import/util_test.go index 3a4e2072c5..80bc31778c 100644 --- a/tests/integration/import/util_test.go +++ b/tests/integration/import/util_test.go @@ -174,6 +174,17 @@ func GenerateNumpyFiles(cm storage.ChunkManager, schema *schemapb.CollectionSche data = chunkedRows case schemapb.DataType_SparseFloatVector: data = insertData.Data[fieldID].(*storage.SparseFloatVectorFieldData).GetContents() + case schemapb.DataType_Int8Vector: + rows := insertData.Data[fieldID].GetDataRows().([]int8) + if dim != fieldData.(*storage.Int8VectorFieldData).Dim { + panic(fmt.Sprintf("dim mis-match: %d, %d", dim, fieldData.(*storage.Int8VectorFieldData).Dim)) + } + chunked := lo.Chunk(rows, dim) + chunkedRows := make([][dim]int8, len(chunked)) + for i, innerSlice := range chunked { + copy(chunkedRows[i][:], innerSlice) + } + data = chunkedRows default: data = insertData.Data[fieldID].GetDataRows() } diff --git a/tests/integration/util_insert.go b/tests/integration/util_insert.go index cf227ea989..1a37538211 100644 --- a/tests/integration/util_insert.go +++ b/tests/integration/util_insert.go @@ -155,6 +155,10 @@ func NewSparseFloatVectorFieldData(fieldName string, numRows int) *schemapb.Fiel return testutils.NewSparseFloatVectorFieldData(fieldName, numRows) } +func NewInt8VectorFieldData(fieldName string, numRows, dim int) *schemapb.FieldData { + return testutils.NewInt8VectorFieldData(fieldName, numRows, dim) +} + func GenerateInt64Array(numRows int, start int64) []int64 { ret := make([]int64, numRows) for i := 0; i < numRows; i++ { diff --git a/tests/integration/util_query.go b/tests/integration/util_query.go index 500b4d34ac..fd14f0099e 100644 --- a/tests/integration/util_query.go +++ b/tests/integration/util_query.go @@ -34,6 +34,7 @@ import ( "github.com/milvus-io/milvus/pkg/util/merr" "github.com/milvus-io/milvus/pkg/util/metricsinfo" "github.com/milvus-io/milvus/pkg/util/testutils" + "github.com/milvus-io/milvus/pkg/util/typeutil" ) const ( @@ -328,6 +329,13 @@ func constructPlaceholderGroup(nq, dim int, vectorType schemapb.DataType) *commo placeholderType = commonpb.PlaceholderType_SparseFloatVector sparseVecs := GenerateSparseFloatArray(nq) values = append(values, sparseVecs.Contents...) + case schemapb.DataType_Int8Vector: + placeholderType = commonpb.PlaceholderType_Int8Vector + data := testutils.GenerateInt8Vectors(nq, dim) + for i := 0; i < nq; i++ { + rowBytes := dim + values = append(values, typeutil.Int8ArrayToBytes(data[rowBytes*i:rowBytes*(i+1)])) + } default: panic("invalid vector data type") }