mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
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:
@@ -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,
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user