fix: heap buffer overflow when unittest at index wrapper (#36838)

issue: #35852

---------

Signed-off-by: chyezh <chyezh@outlook.com>
pull/36859/head
Zhen Ye 2024-10-14 18:13:22 +08:00 committed by GitHub
parent 5ec4163d0f
commit f46c3acea9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 5 deletions

View File

@ -172,9 +172,9 @@ TEST_P(IndexWrapperTest, BuildAndQuery) {
search_info.search_params_ = search_conf;
std::unique_ptr<SearchResult> result;
if (vec_field_data_type == DataType::VECTOR_FLOAT) {
auto dataset = GenDataset(NB, metric_type, false);
auto nb_for_nq = NQ + query_offset;
auto dataset = GenDataset(nb_for_nq, metric_type, false);
auto xb_data = dataset.get_col<float>(milvus::FieldId(100));
auto xb_dataset = knowhere::GenDataSet(NB, DIM, xb_data.data());
auto xq_dataset =
knowhere::GenDataSet(NQ, DIM, xb_data.data() + DIM * query_offset);
result = vec_index->Query(xq_dataset, search_info, nullptr);
@ -188,11 +188,12 @@ TEST_P(IndexWrapperTest, BuildAndQuery) {
xq_dataset->SetIsSparse(true);
result = vec_index->Query(xq_dataset, search_info, nullptr);
} else {
auto dataset = GenDataset(NB, metric_type, true);
auto nb_for_nq = NQ + query_offset;
auto dataset = GenDataset(nb_for_nq, metric_type, true);
auto xb_bin_data = dataset.get_col<uint8_t>(milvus::FieldId(100));
auto xb_dataset = knowhere::GenDataSet(NB, DIM, xb_bin_data.data());
// offset of binary vector is 8-aligned bit-wise representation.
auto xq_dataset = knowhere::GenDataSet(
NQ, DIM, xb_bin_data.data() + DIM * query_offset);
NQ, DIM, xb_bin_data.data() + ((DIM + 7) / 8) * query_offset);
result = vec_index->Query(xq_dataset, search_info, nullptr);
}