fix: fix corruption of offsets between milvus and knowhere (#48754)

issue: https://github.com/milvus-io/milvus/issues/48743

---------

Signed-off-by: SpadeA-Tang <tangchenjie1210@gmail.com>
Signed-off-by: SpadeA <tangchenjie1210@gmail.com>
This commit is contained in:
Spade A
2026-04-07 21:05:42 +08:00
committed by GitHub
parent 86bdb30956
commit e7e321c0a1
3 changed files with 320 additions and 24 deletions
+17 -8
View File
@@ -285,6 +285,10 @@ VectorDiskAnnIndex<T>::BuildWithDataset(const DatasetPtr& dataset,
storage::LocalChunkManagerSingleton::GetInstance().GetChunkManager();
knowhere::Json build_config;
build_config.update(config);
auto is_embedding_list = (elem_type_ != DataType::NONE);
build_config[EMB_LIST] = is_embedding_list;
// set data path
auto segment_id = file_manager_->GetFieldDataMeta().segment_id;
auto field_id = file_manager_->GetFieldDataMeta().field_id;
@@ -338,18 +342,23 @@ VectorDiskAnnIndex<T>::BuildWithDataset(const DatasetPtr& dataset,
"offset";
local_chunk_manager->CreateFile(offsets_path);
// Calculate the number of offsets (num_rows + 1)
// We need to find the actual number by looking at the data
uint32_t num_rows =
static_cast<uint32_t>(milvus::GetDatasetRows(dataset));
uint32_t num_offsets = num_rows + 1;
// GetDatasetRows returns total flattened vector count for vector arrays,
// not the number of emb_lists. Count actual offsets by scanning the array
// until we reach the terminal element (== total_vectors).
size_t total_vectors =
static_cast<size_t>(milvus::GetDatasetRows(dataset));
size_t num_offsets = 0;
while (offsets[num_offsets] < total_vectors) {
num_offsets++;
}
num_offsets++; // include the terminal element (== total_vectors)
// Write offsets to file
// Format: [num_offsets][offsets_data]
// Format: [num_offsets (size_t)][offsets_data (size_t array)]
int64_t write_pos = 0;
local_chunk_manager->Write(
offsets_path, write_pos, &num_offsets, sizeof(uint32_t));
write_pos += sizeof(uint32_t);
offsets_path, write_pos, &num_offsets, sizeof(size_t));
write_pos += sizeof(size_t);
local_chunk_manager->Write(
offsets_path,
@@ -565,6 +565,12 @@ DiskFileManagerImpl::cache_raw_data_to_disk_internal(const Config& config) {
FetchRawData();
}
// For vector arrays, num_rows should be the total flattened vector count,
// not the number of emb_lists, because DiskANN reads this from the data file header.
if (is_vector_array) {
num_rows = static_cast<uint32_t>(offsets.back());
}
// write num_rows and dim value to file header
write_offset = 0;
local_chunk_manager->Write(
@@ -575,23 +581,23 @@ DiskFileManagerImpl::cache_raw_data_to_disk_internal(const Config& config) {
// Write offsets file for VECTOR_ARRAY
if (is_vector_array) {
AssertInfo(offsets.size() == num_rows + 1,
"offsets size is not equal to num_rows + 1: offset size {}, "
"num_rows {}",
offsets.size(),
num_rows);
AssertInfo(
offsets.size() >= 2 && offsets.front() == 0,
"invalid emb_list offsets: size {}, front {}",
offsets.size(),
offsets.empty() ? -1 : static_cast<int64_t>(offsets.front()));
// Get offsets path from config if provided, otherwise use default
auto offsets_path = index::GetValueFromConfig<std::string>(
config, index::EMB_LIST_OFFSETS_PATH)
.value();
local_chunk_manager->CreateFile(offsets_path);
uint32_t num_offsets = offsets.size();
size_t num_offsets = offsets.size();
int64_t offsets_write_pos = 0;
local_chunk_manager->Write(
offsets_path, offsets_write_pos, &num_offsets, sizeof(uint32_t));
offsets_write_pos += sizeof(uint32_t);
offsets_path, offsets_write_pos, &num_offsets, sizeof(size_t));
offsets_write_pos += sizeof(size_t);
local_chunk_manager->Write(offsets_path,
offsets_write_pos,
@@ -831,6 +837,12 @@ DiskFileManagerImpl::cache_raw_data_to_disk_storage_v2(const Config& config) {
is_vector_array ? &offsets : nullptr);
}
// For vector arrays, num_rows should be the total flattened vector count,
// not the number of emb_lists, because DiskANN reads this from the data file header.
if (is_vector_array) {
num_rows = static_cast<uint32_t>(offsets.back());
}
// write num_rows and dim value to file header
write_offset = 0;
local_chunk_manager->Write(
@@ -841,11 +853,11 @@ DiskFileManagerImpl::cache_raw_data_to_disk_storage_v2(const Config& config) {
// Write offsets file for VECTOR_ARRAY
if (is_vector_array) {
AssertInfo(offsets.size() == num_rows + 1,
"offsets size is not equal to num_rows + 1: offset size {}, "
"num_rows {}",
offsets.size(),
num_rows);
AssertInfo(
offsets.size() >= 2 && offsets.front() == 0,
"invalid emb_list offsets: size {}, front {}",
offsets.size(),
offsets.empty() ? -1 : static_cast<int64_t>(offsets.front()));
// Get offsets path from config if provided, otherwise use default
auto offsets_path = index::GetValueFromConfig<std::string>(
config, index::EMB_LIST_OFFSETS_PATH)
@@ -853,12 +865,12 @@ DiskFileManagerImpl::cache_raw_data_to_disk_storage_v2(const Config& config) {
local_chunk_manager->CreateFile(offsets_path);
uint32_t num_offsets = offsets.size();
size_t num_offsets = offsets.size();
int64_t offsets_write_pos = 0;
local_chunk_manager->Write(
offsets_path, offsets_write_pos, &num_offsets, sizeof(uint32_t));
offsets_write_pos += sizeof(uint32_t);
offsets_path, offsets_write_pos, &num_offsets, sizeof(size_t));
offsets_write_pos += sizeof(size_t);
local_chunk_manager->Write(offsets_path,
offsets_write_pos,
+275
View File
@@ -9,6 +9,7 @@
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License
#include <boost/filesystem/operations.hpp>
#include <folly/FBVector.h>
#include <gtest/gtest.h>
#include <nlohmann/json_fwd.hpp>
@@ -29,11 +30,13 @@
#include "bitset/bitset.h"
#include "common/BitsetView.h"
#include "common/Consts.h"
#include "common/EasyAssert.h"
#include "common/QueryInfo.h"
#include "common/QueryResult.h"
#include "common/Schema.h"
#include "common/Tracer.h"
#include "common/VectorArray.h"
#include "common/TypeTraits.h"
#include "common/Types.h"
#include "common/protobuf_utils.h"
@@ -62,6 +65,8 @@
#include "segcore/ReduceStructure.h"
#include "segcore/reduce/Reduce.h"
#include "storage/FileManager.h"
#include "storage/InsertData.h"
#include "storage/PayloadReader.h"
#include "storage/ThreadPools.h"
#include "storage/Types.h"
#include "storage/Util.h"
@@ -1059,6 +1064,276 @@ TEST(Indexing, SearchDiskAnnWithBFloat16) {
EXPECT_NO_THROW(
vec_index->Query(xq_dataset, search_info, nullptr, nullptr, result));
}
TEST(Indexing, DiskAnnEmbListBuildWithDataset) {
// Test BuildWithDataset path for VECTOR_ARRAY (embedding list) with DiskANN.
// Each emb_list has a different number of vectors to exercise the offset logic.
int64_t num_emb_lists = 100;
int64_t K = 4;
IndexType index_type = knowhere::IndexEnum::INDEX_DISKANN;
MetricType metric_type = knowhere::metric::MAX_SIM_L2;
// Variable-length emb_lists: lengths cycle through 1,2,3,4,5
std::vector<size_t> offsets;
offsets.push_back(0);
for (int64_t i = 0; i < num_emb_lists; ++i) {
size_t len = (i % 5) + 1; // 1..5 vectors per emb_list
offsets.push_back(offsets.back() + len);
}
int64_t total_vectors = static_cast<int64_t>(offsets.back());
// Generate random float vectors
std::mt19937 rng(42);
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
std::vector<float> xb_data(total_vectors * DIM);
for (auto& v : xb_data) {
v = dist(rng);
}
// Create dataset with EMB_LIST_OFFSET
knowhere::DataSetPtr xb_dataset =
knowhere::GenDataSet(total_vectors, DIM, xb_data.data());
xb_dataset->Set(knowhere::meta::EMB_LIST_OFFSET,
const_cast<const size_t*>(offsets.data()));
// Setup file manager context
int64_t collection_id = 1;
int64_t partition_id = 2;
int64_t segment_id = 3;
int64_t field_id = 100;
int64_t build_id = 1000;
int64_t index_version = 1;
StorageConfig storage_config = get_default_local_storage_config();
milvus::storage::FieldDataMeta field_data_meta{
collection_id, partition_id, segment_id, field_id};
field_data_meta.field_schema.set_data_type(
static_cast<proto::schema::DataType>(milvus::DataType::VECTOR_ARRAY));
field_data_meta.field_schema.set_element_type(
static_cast<proto::schema::DataType>(milvus::DataType::VECTOR_FLOAT));
milvus::storage::IndexMeta index_meta{
segment_id, field_id, build_id, index_version};
auto chunk_manager = storage::CreateChunkManager(storage_config);
auto fs = milvus::storage::InitArrowFileSystem(storage_config);
milvus::storage::FileManagerContext file_manager_context(
field_data_meta, index_meta, chunk_manager, fs);
// Create index with elem_type = VECTOR_FLOAT (embedding list mode)
milvus::index::CreateIndexInfo create_index_info;
create_index_info.index_type = index_type;
create_index_info.metric_type = metric_type;
create_index_info.field_type = milvus::DataType::VECTOR_ARRAY;
create_index_info.index_engine_version =
knowhere::Version::GetCurrentVersion().VersionNumber();
auto index = milvus::index::IndexFactory::GetInstance().CreateIndex(
create_index_info, file_manager_context);
auto build_conf = Config{
{knowhere::meta::METRIC_TYPE, metric_type},
{knowhere::meta::DIM, std::to_string(DIM)},
{milvus::index::DISK_ANN_MAX_DEGREE, std::to_string(24)},
{milvus::index::DISK_ANN_SEARCH_LIST_SIZE, std::to_string(56)},
{milvus::index::DISK_ANN_PQ_CODE_BUDGET, std::to_string(0.001)},
{milvus::index::DISK_ANN_BUILD_DRAM_BUDGET, std::to_string(2)},
{milvus::index::DISK_ANN_BUILD_THREAD_NUM, std::to_string(2)},
};
// Build
ASSERT_NO_THROW(index->BuildWithDataset(xb_dataset, build_conf));
// Upload -> Load
auto create_index_result = index->Upload();
ASSERT_GT(create_index_result->GetMemSize(), 0);
ASSERT_GT(create_index_result->GetSerializedSize(), 0);
auto index_files = create_index_result->GetIndexFiles();
index.reset();
auto new_index = milvus::index::IndexFactory::GetInstance().CreateIndex(
create_index_info, file_manager_context);
auto vec_index = dynamic_cast<milvus::index::VectorIndex*>(new_index.get());
auto load_conf = generate_load_conf(index_type, metric_type, total_vectors);
load_conf["index_files"] = index_files;
load_conf[milvus::LOAD_PRIORITY] =
milvus::proto::common::LoadPriority::HIGH;
vec_index->Load(milvus::tracer::TraceContext{}, load_conf);
EXPECT_EQ(vec_index->Count(), total_vectors);
// Search: use 2 query emb_lists (3 vectors + 2 vectors)
int64_t NQ_vecs = 5;
std::vector<float> query_data(NQ_vecs * DIM);
for (auto& v : query_data) {
v = dist(rng);
}
knowhere::DataSetPtr xq_dataset =
knowhere::GenDataSet(NQ_vecs, DIM, query_data.data());
std::vector<size_t> query_offsets = {0, 3, 5};
xq_dataset->Set(knowhere::meta::EMB_LIST_OFFSET,
const_cast<const size_t*>(query_offsets.data()));
milvus::SearchInfo search_info;
search_info.topk_ = K;
search_info.metric_type_ = metric_type;
search_info.search_params_ = milvus::Config{
{knowhere::meta::METRIC_TYPE, metric_type},
{milvus::index::DISK_ANN_QUERY_LIST, K * 2},
};
SearchResult result;
EXPECT_NO_THROW(
vec_index->Query(xq_dataset, search_info, nullptr, nullptr, result));
// 2 query emb_lists
EXPECT_EQ(result.total_nq_, 2);
EXPECT_EQ(result.distances_.size(), 2 * K);
// Cleanup local index data
vec_index->CleanLocalData();
}
TEST(Indexing, DiskAnnEmbListBuildFromBinlog) {
// Test Build path (from binlog) for VECTOR_ARRAY with DiskANN.
int64_t num_emb_lists = 100;
int64_t K = 4;
int64_t dim = DIM;
IndexType index_type = knowhere::IndexEnum::INDEX_DISKANN;
MetricType metric_type = knowhere::metric::MAX_SIM_L2;
int64_t collection_id = 1;
int64_t partition_id = 2;
int64_t segment_id = 3;
int64_t field_id = 100;
int64_t build_id = 1000;
int64_t index_version = 1;
// Generate VECTOR_ARRAY field data with variable-length emb_lists
std::mt19937 rng(42);
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
auto field_data = storage::CreateFieldData(
DataType::VECTOR_ARRAY, DataType::VECTOR_FLOAT, false, dim);
std::vector<milvus::VectorArray> vec_arrays;
int64_t total_vectors = 0;
for (int64_t i = 0; i < num_emb_lists; ++i) {
int num_vecs = static_cast<int>((i % 5) + 1); // 1..5 vectors
std::vector<float> vecs(num_vecs * dim);
for (auto& v : vecs) {
v = dist(rng);
}
vec_arrays.emplace_back(
vecs.data(), num_vecs, dim, DataType::VECTOR_FLOAT);
total_vectors += num_vecs;
}
field_data->FillFieldData(vec_arrays.data(), vec_arrays.size());
// Serialize to binlog
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
storage::InsertData insert_data(payload_reader);
auto field_data_meta =
milvus::segcore::gen_field_meta(collection_id,
partition_id,
segment_id,
field_id,
DataType::VECTOR_ARRAY,
DataType::VECTOR_FLOAT,
false);
insert_data.SetFieldDataMeta(field_data_meta);
insert_data.SetTimestamps(0, 100);
auto serialized_bytes = insert_data.Serialize(storage::StorageType::Remote);
// Write binlog to storage
StorageConfig storage_config = get_default_local_storage_config();
auto chunk_manager = storage::CreateChunkManager(storage_config);
auto fs = milvus::storage::InitArrowFileSystem(storage_config);
std::string log_root = fmt::format("{}{}", TestRemotePath, collection_id);
boost::filesystem::remove_all(log_root);
std::string log_dir = fmt::format(
"{}/{}/{}/{}", log_root, partition_id, segment_id, field_id);
std::string log_path = log_dir + "/0";
chunk_manager->Write(
log_path, serialized_bytes.data(), serialized_bytes.size());
// Create index via Build path
milvus::storage::IndexMeta index_meta{segment_id,
field_id,
build_id,
index_version,
"test",
"vec_field",
DataType::VECTOR_FLOAT,
dim};
milvus::storage::FileManagerContext file_manager_context(
field_data_meta, index_meta, chunk_manager, fs);
milvus::index::CreateIndexInfo create_index_info;
create_index_info.index_type = index_type;
create_index_info.metric_type = metric_type;
create_index_info.field_type = milvus::DataType::VECTOR_ARRAY;
create_index_info.index_engine_version =
knowhere::Version::GetCurrentVersion().VersionNumber();
auto index = milvus::index::IndexFactory::GetInstance().CreateIndex(
create_index_info, file_manager_context);
Config build_conf;
build_conf[knowhere::meta::METRIC_TYPE] = metric_type;
build_conf[knowhere::meta::DIM] = std::to_string(dim);
build_conf[milvus::index::DISK_ANN_MAX_DEGREE] = std::to_string(24);
build_conf[milvus::index::DISK_ANN_SEARCH_LIST_SIZE] = std::to_string(56);
build_conf[milvus::index::DISK_ANN_PQ_CODE_BUDGET] = std::to_string(0.001);
build_conf[milvus::index::DISK_ANN_BUILD_DRAM_BUDGET] = std::to_string(2);
build_conf[milvus::index::DISK_ANN_BUILD_THREAD_NUM] = std::to_string(2);
build_conf[INSERT_FILES_KEY] = std::vector<std::string>{log_path};
ASSERT_NO_THROW(index->Build(build_conf));
// Upload -> Load
auto create_index_result = index->Upload();
ASSERT_GT(create_index_result->GetMemSize(), 0);
ASSERT_GT(create_index_result->GetSerializedSize(), 0);
auto index_files = create_index_result->GetIndexFiles();
index.reset();
auto new_index = milvus::index::IndexFactory::GetInstance().CreateIndex(
create_index_info, file_manager_context);
auto vec_index = dynamic_cast<milvus::index::VectorIndex*>(new_index.get());
auto load_conf = generate_load_conf(index_type, metric_type, total_vectors);
load_conf["index_files"] = index_files;
load_conf[milvus::LOAD_PRIORITY] =
milvus::proto::common::LoadPriority::HIGH;
vec_index->Load(milvus::tracer::TraceContext{}, load_conf);
EXPECT_EQ(vec_index->Count(), total_vectors);
// Search: 2 query emb_lists (3 vectors + 2 vectors)
int64_t NQ_vecs = 5;
std::vector<float> query_data(NQ_vecs * dim);
for (auto& v : query_data) {
v = dist(rng);
}
knowhere::DataSetPtr xq_dataset =
knowhere::GenDataSet(NQ_vecs, dim, query_data.data());
std::vector<size_t> query_offsets = {0, 3, 5};
xq_dataset->Set(knowhere::meta::EMB_LIST_OFFSET,
const_cast<const size_t*>(query_offsets.data()));
milvus::SearchInfo search_info;
search_info.topk_ = K;
search_info.metric_type_ = metric_type;
search_info.search_params_ = milvus::Config{
{knowhere::meta::METRIC_TYPE, metric_type},
{milvus::index::DISK_ANN_QUERY_LIST, K * 2},
};
SearchResult result;
EXPECT_NO_THROW(
vec_index->Query(xq_dataset, search_info, nullptr, nullptr, result));
EXPECT_EQ(result.total_nq_, 2);
EXPECT_EQ(result.distances_.size(), 2 * K);
// Cleanup: local index data + binlog
vec_index->CleanLocalData();
boost::filesystem::remove_all(log_root);
}
#endif
TEST(Indexing, IndexStats) {