diff --git a/internal/core/src/index/VectorDiskIndex.cpp b/internal/core/src/index/VectorDiskIndex.cpp index 0de13e20ae..7e74097e37 100644 --- a/internal/core/src/index/VectorDiskIndex.cpp +++ b/internal/core/src/index/VectorDiskIndex.cpp @@ -285,6 +285,10 @@ VectorDiskAnnIndex::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::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(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(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, diff --git a/internal/core/src/storage/DiskFileManagerImpl.cpp b/internal/core/src/storage/DiskFileManagerImpl.cpp index 65591c69eb..16ebe337f5 100644 --- a/internal/core/src/storage/DiskFileManagerImpl.cpp +++ b/internal/core/src/storage/DiskFileManagerImpl.cpp @@ -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(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(offsets.front())); // Get offsets path from config if provided, otherwise use default auto offsets_path = index::GetValueFromConfig( 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(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(offsets.front())); // Get offsets path from config if provided, otherwise use default auto offsets_path = index::GetValueFromConfig( 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, diff --git a/internal/core/unittest/test_indexing.cpp b/internal/core/unittest/test_indexing.cpp index 9503f356f4..435ed04b59 100644 --- a/internal/core/unittest/test_indexing.cpp +++ b/internal/core/unittest/test_indexing.cpp @@ -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 #include #include #include @@ -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 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(offsets.back()); + + // Generate random float vectors + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::vector 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(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(milvus::DataType::VECTOR_ARRAY)); + field_data_meta.field_schema.set_element_type( + static_cast(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(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 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 query_offsets = {0, 3, 5}; + xq_dataset->Set(knowhere::meta::EMB_LIST_OFFSET, + const_cast(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 dist(-1.0f, 1.0f); + + auto field_data = storage::CreateFieldData( + DataType::VECTOR_ARRAY, DataType::VECTOR_FLOAT, false, dim); + + std::vector vec_arrays; + int64_t total_vectors = 0; + for (int64_t i = 0; i < num_emb_lists; ++i) { + int num_vecs = static_cast((i % 5) + 1); // 1..5 vectors + std::vector 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(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{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(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 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 query_offsets = {0, 3, 5}; + xq_dataset->Set(knowhere::meta::EMB_LIST_OFFSET, + const_cast(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) {