diff --git a/include/utils.h b/include/utils.h index beb7d31205..9635e0ab85 100644 --- a/include/utils.h +++ b/include/utils.h @@ -3,7 +3,10 @@ #pragma once +#include #include +#include +#include #include "common_includes.h" @@ -182,11 +185,69 @@ inline void convert_labels_string_to_int(const std::string &inFileName, const st const std::string &mapFileName, const std::string &unv_label, uint32_t& unv_label_id) { + // Optimized version of the original per-line stdio implementation. + // + // Hotspots eliminated: + // 1. `label_writer << ... << std::endl` (88M flushes → ~1700s): + // replaced with a batched writer that flushes in 32 MB chunks. + // 2. `label_writer << lbls[j]` per int (stdio format): replaced with + // locale-independent std::to_chars conversion. + // + // Byte-for-byte output preservation: + // - Label IDs are assigned in the same "first seen" order, so the + // unordered_map contents are identical to the original. + // - Row order in the output file matches the input, same as original. + // - Every row is terminated by std::endl on Windows the ofstream would + // translate '\n' to "\r\n" when we used `<<`; the optimized writer + // uses binary mode and emits CRLF + // explicitly on Windows to keep the output bytes identical. + // - The map file is written from an unordered_map, whose iteration + // order was already unspecified in the original code; we keep the + // same iteration, so semantically both are equivalent for + // downstream load_label_map (which does not rely on order). + std::unordered_map string_int_map; - std::ofstream label_writer(outFileName); + + // ---- Writer for the formatted labels file ---- + std::ofstream label_writer(outFileName, std::ios::binary); + if (!label_writer.is_open()) + { + throw diskann::ANNException(std::string("Failed to open ") + outFileName, -1); + } + + // Bump the stdio buffer as a safety net in case any write bypasses our batch. + constexpr size_t kStdioBufSize = 16 * 1024 * 1024; // 16 MB + std::vector stdio_buf(kStdioBufSize); + label_writer.rdbuf()->pubsetbuf(stdio_buf.data(), stdio_buf.size()); + + constexpr size_t kBatchSize = 32 * 1024 * 1024; // 32 MB + std::vector batch; + batch.reserve(kBatchSize + 4096); // small slack for line-worst-case + + auto flush_batch = [&]() { + if (!batch.empty()) + { + label_writer.write(batch.data(), batch.size()); + batch.clear(); + } + }; + + auto append_uint32 = [](std::vector &dst, uint32_t v) { + char buffer[10]; // uint32 max: 10 digits + const auto result = std::to_chars(buffer, buffer + sizeof(buffer), v); + if (result.ec != std::errc()) + { + throw diskann::ANNException("Failed to format label ID", -1); + } + dst.insert(dst.end(), buffer, result.ptr); + }; + std::ifstream label_reader(inFileName); - //if (unv_label != "") - // string_int_map[unv_label] = 0; + if (!label_reader.is_open()) + { + throw diskann::ANNException(std::string("Failed to open ") + inFileName, -1); + } + std::string line, token; while (std::getline(label_reader, line)) { @@ -208,14 +269,28 @@ inline void convert_labels_string_to_int(const std::string &inFileName, const st std::cout << "No label found"; exit(-1); } + + // Emit "id1,id2,...,idN\n" (with CRLF on Windows to preserve byte + // equivalence with the original text-mode ofstream + std::endl). for (size_t j = 0; j < lbls.size(); j++) { - if (j != lbls.size() - 1) - label_writer << lbls[j] << ","; - else - label_writer << lbls[j] << std::endl; + append_uint32(batch, lbls[j]); + if (j + 1 < lbls.size()) + { + batch.push_back(','); + } + } +#ifdef _WIN32 + batch.push_back('\r'); +#endif + batch.push_back('\n'); + + if (batch.size() >= kBatchSize) + { + flush_batch(); } } + flush_batch(); label_writer.close(); if (unv_label != "") @@ -228,8 +303,13 @@ inline void convert_labels_string_to_int(const std::string &inFileName, const st // else: unv_label_id remains 0, indicating label not found } + // ---- Writer for the map file. 235 or so entries — tiny, keep simple. ---- std::ofstream map_writer(mapFileName); - for (auto mp : string_int_map) + if (!map_writer.is_open()) + { + throw diskann::ANNException(std::string("Failed to open ") + mapFileName, -1); + } + for (auto &mp : string_int_map) { map_writer << mp.first << "\t" << mp.second << std::endl; } @@ -671,16 +751,45 @@ inline void load_bin(MemoryMappedFiles &files, const std::string &bin_file, std: inline void copy_file(std::string in_file, std::string out_file) { - std::ifstream source(in_file, std::ios::binary); - std::ofstream dest(out_file, std::ios::binary); - - std::istreambuf_iterator begin_source(source); - std::istreambuf_iterator end_source; - std::ostreambuf_iterator begin_dest(dest); - std::copy(begin_source, end_source, begin_dest); + // OS-level fast copy. Replaces the previous istreambuf_iterator + // per-character copy which was ~1 MB/s on 200 MB label files. + // std::filesystem::copy_file uses the OS's fast copy primitive + // (CopyFileW on Windows, copy_file_range/sendfile on Linux). + std::filesystem::copy_file( + in_file, out_file, + std::filesystem::copy_options::overwrite_existing); +} + +// move_file — rename-based fast move with copy+remove fallback. +// +// Use this instead of `copy_file(src, dst); std::remove(src);` when the +// caller's intent is to *move* the file to a new location. In practice +// this is O(1) on the same filesystem (~0 ms) versus ~1 s per 200 MB +// for the copy path. +// +// Falls back to copy+remove if rename fails, which can happen for: +// - cross-filesystem/cross-device moves (rename is filesystem-local) +// - the source being held open by another process (antivirus, etc.) +// - unusual permission scenarios +// +// If the target already exists, it is removed first (best-effort) so +// that rename can succeed on platforms where rename refuses to overwrite. +inline void move_file(std::string in_file, std::string out_file) +{ + std::error_code ec; + std::filesystem::remove(out_file, ec); // best-effort clear target; ignore errors + ec.clear(); + std::filesystem::rename(in_file, out_file, ec); + if (!ec) + { + return; // fast path succeeded + } - source.close(); - dest.close(); + // Fallback: copy + remove. + std::filesystem::copy_file( + in_file, out_file, + std::filesystem::copy_options::overwrite_existing); + std::filesystem::remove(in_file, ec); // best-effort; ignore errors } DISKANN_DLLEXPORT double calculate_recall(unsigned num_queries, unsigned *gold_std, float *gs_dist, unsigned dim_gs, diff --git a/src/disk_utils.cpp b/src/disk_utils.cpp index 40b0a7c134..ac2024b556 100644 --- a/src/disk_utils.cpp +++ b/src/disk_utils.cpp @@ -698,12 +698,11 @@ int build_merged_vamana_index(std::string base_file, diskann::Metric compareMetr if (use_filters) { - // need to copy the labels_to_medoids file to the specified input + // need to move the labels_to_medoids file to the specified input // file std::remove(labels_to_medoids_file.c_str()); std::string mem_labels_to_medoid_file = mem_index_path + "_labels_to_medoids.txt"; - copy_file(mem_labels_to_medoid_file, labels_to_medoids_file); - std::remove(mem_labels_to_medoid_file.c_str()); + move_file(mem_labels_to_medoid_file, labels_to_medoids_file); } std::remove(medoids_file.c_str()); @@ -1396,46 +1395,51 @@ int build_disk_index(const char *dataFilePath, const char *indexFilePath, const gen_random_slice(data_file_to_use.c_str(), sample_base_prefix, sample_sampling_rate); if (use_filters) { - copy_file(labels_file_to_use, disk_labels_file); + // Move the formatted labels file into its final disk-index location. + // The original code did copy_file + std::remove of a *different* + // mem_labels_file (the one Index::save wrote separately). Keep the + // remove of mem_labels_file (that file is redundant with + // labels_file_to_use in the static index case) and move + // labels_file_to_use into place. + move_file(labels_file_to_use, disk_labels_file); std::remove(mem_labels_file.c_str()); if (universal_label != "" && universal_label_id != 0) { - copy_file(mem_univ_label_file, disk_univ_label_file); - std::remove(mem_univ_label_file.c_str()); + move_file(mem_univ_label_file, disk_univ_label_file); } // rename bimask label file std::string bitmask_label_file = std::string(mem_index_path) + "_bitmask_labels.bin"; if (file_exists(bitmask_label_file)) { - copy_file(bitmask_label_file, disk_bitmask_labels_file); - std::remove(bitmask_label_file.c_str()); + move_file(bitmask_label_file, disk_bitmask_labels_file); } - + // rename integer label file std::string integer_label_file = std::string(mem_index_path) + "_integer_labels.bin"; if (file_exists(integer_label_file)) { - copy_file(integer_label_file, disk_integer_labels_file); - std::remove(integer_label_file.c_str()); + move_file(integer_label_file, disk_integer_labels_file); } std::remove(augmented_data_file.c_str()); std::remove(augmented_labels_file.c_str()); + // labels_file_to_use was already moved to disk_labels_file above; + // the remove below is a no-op if that succeeded, and cleans up + // the copy-fallback case where move_file left the source behind + // because of a permission or cross-filesystem failure. std::remove(labels_file_to_use.c_str()); } - + std::string old_seller_mem_file = std::string(mem_index_path) + "_sellers.txt"; if (file_exists(old_seller_mem_file)) { - copy_file(old_seller_mem_file, old_disk_seller_file); - std::remove(old_seller_mem_file.c_str()); + move_file(old_seller_mem_file, old_disk_seller_file); } std::string seller_mem_file = std::string(mem_index_path) + "_sellers.bin"; if (file_exists(seller_mem_file)) { - copy_file(seller_mem_file, disk_seller_file); - std::remove(seller_mem_file.c_str()); + move_file(seller_mem_file, disk_seller_file); } if (created_temp_file_for_processed_data) @@ -1443,7 +1447,6 @@ int build_disk_index(const char *dataFilePath, const char *indexFilePath, const std::remove(mem_index_path.c_str()); if (use_disk_pq) std::remove(disk_pq_compressed_vectors_path.c_str()); - auto e = std::chrono::high_resolution_clock::now(); std::chrono::duration diff = e - s; diskann::cout << "Indexing time: " << diff.count() << std::endl; diff --git a/src/in_mem_graph_store.cpp b/src/in_mem_graph_store.cpp index fbd9ca750f..0f1e95055c 100644 --- a/src/in_mem_graph_store.cpp +++ b/src/in_mem_graph_store.cpp @@ -4,6 +4,9 @@ #include "in_mem_graph_store.h" #include "utils.h" +#include +#include + namespace diskann { @@ -204,32 +207,82 @@ std::tuple InMemGraphStore::load_impl(const std::str int InMemGraphStore::save_graph(const std::string &index_path_prefix, const size_t num_points, const size_t num_frozen_points, const uint32_t start) { + // Save graph to disk in batches to avoid the massive syscall overhead of + // the original per-node ofstream::write() calls. + // + // Original: 2 writes/node * ~88M nodes = ~176M syscalls, + // observed throughput ~11 MB/s on 18.7 GB output. + // New: Serialize into 32 MB batches; ~600 syscalls total. + // + // Byte-for-byte equivalent to the original layout: + // header (24B): [index_size u64][max_observed_degree u32][ep u32][num_frozen u64] + // per node: [GK u32][neighbors[GK] u32] + // header is back-patched at the end with the true index_size + max_degree. + std::ofstream out; open_file_to_write(out, index_path_prefix); + // Bump the stdio buffer as a safety net in case any write bypasses our batch. + constexpr size_t kStdioBufSize = 16 * 1024 * 1024; // 16 MB + std::vector stdio_buf(kStdioBufSize); + out.rdbuf()->pubsetbuf(stdio_buf.data(), stdio_buf.size()); + size_t file_offset = 0; out.seekp(file_offset, out.beg); size_t index_size = 24; uint32_t max_degree = 0; + + // Placeholder header (overwritten at the end). out.write((char *)&index_size, sizeof(uint64_t)); out.write((char *)&_max_observed_degree, sizeof(uint32_t)); uint32_t ep_u32 = start; out.write((char *)&ep_u32, sizeof(uint32_t)); out.write((char *)&num_frozen_points, sizeof(size_t)); - // Note: num_points = _nd + _num_frozen_points + // Batched write: accumulate 32 MB of node data before hitting the disk. + constexpr size_t kBatchSize = 32 * 1024 * 1024; // 32 MB + // Worst-case node bytes: (1 header uint32 + max_degree neighbors) * 4 B. + // Use a generous cap so we never overshoot the batch by too much. + const size_t kMaxNodeBytes = (static_cast(_max_observed_degree) + 4) * sizeof(uint32_t); + std::vector batch; + batch.reserve(kBatchSize + kMaxNodeBytes); + + auto flush_batch = [&]() { + if (!batch.empty()) + { + out.write(batch.data(), batch.size()); + batch.clear(); + } + }; + for (uint32_t i = 0; i < num_points; i++) { uint32_t GK = (uint32_t)_graph[i].size(); - out.write((char *)&GK, sizeof(uint32_t)); - out.write((char *)_graph[i].data(), GK * sizeof(uint32_t)); + size_t node_bytes = sizeof(uint32_t) * (1 + GK); + + // Append GK followed by the neighbor array to the batch buffer. + size_t old_size = batch.size(); + batch.resize(old_size + node_bytes); + std::memcpy(batch.data() + old_size, &GK, sizeof(uint32_t)); + std::memcpy(batch.data() + old_size + sizeof(uint32_t), + _graph[i].data(), GK * sizeof(uint32_t)); + max_degree = _graph[i].size() > max_degree ? (uint32_t)_graph[i].size() : max_degree; - index_size += (size_t)(sizeof(uint32_t) * (GK + 1)); + index_size += node_bytes; + + if (batch.size() >= kBatchSize) + { + flush_batch(); + } } + flush_batch(); + + // Back-patch the header with the actual index size and max degree. out.seekp(file_offset, out.beg); out.write((char *)&index_size, sizeof(uint64_t)); out.write((char *)&max_degree, sizeof(uint32_t)); out.close(); + return (int)index_size; } diff --git a/src/index.cpp b/src/index.cpp index 7433b9fbca..6c03b295ec 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -3,6 +3,7 @@ #include #include +#include #include @@ -336,19 +337,53 @@ void Index::save(const char *filename, bool compact_before_save if (_location_to_labels.size() > 0) { - std::ofstream label_writer(std::string(filename) + "_labels.txt"); + // Avoid the per-line flush caused by std::endl on large label files. + std::ofstream label_writer(std::string(filename) + "_labels.txt", std::ios::binary); assert(label_writer.is_open()); + + constexpr size_t kBatchSize = 32 * 1024 * 1024; // 32 MB + std::vector batch; + batch.reserve(kBatchSize); + + auto flush_batch = [&]() { + if (!batch.empty()) + { + label_writer.write(batch.data(), batch.size()); + batch.clear(); + } + }; + + auto append_label = [](std::vector &dst, LabelT label) { + char buffer[20]; // uint64 max: 20 digits + const auto result = std::to_chars(buffer, buffer + sizeof(buffer), label); + if (result.ec != std::errc()) + { + throw diskann::ANNException("Failed to format label ID", -1); + } + dst.insert(dst.end(), buffer, result.ptr); + }; + for (uint32_t i = 0; i < _nd; i++) { for (uint32_t j = 0; j + 1 < _location_to_labels[i].size(); j++) { - label_writer << _location_to_labels[i][j] << ","; + append_label(batch, _location_to_labels[i][j]); + batch.push_back(','); } if (_location_to_labels[i].size() != 0) - label_writer << _location_to_labels[i][_location_to_labels[i].size() - 1]; + append_label(batch, _location_to_labels[i][_location_to_labels[i].size() - 1]); + +#ifdef _WIN32 + batch.push_back('\r'); +#endif + batch.push_back('\n'); - label_writer << std::endl; + if (batch.size() >= kBatchSize) + { + flush_batch(); + } } + flush_batch(); label_writer.close(); // write compacted raw_labels if data hence _location_to_labels was also compacted @@ -1657,7 +1692,6 @@ void Index::prune_neighbors(const uint32_t location, std::vecto for (auto &ngh : pool) ngh.distance = _data_store->get_distance(ngh.id, location); } - // sort the pool based on distance to query and prune it with occlude_list std::sort(pool.begin(), pool.end()); pruned_list.clear(); @@ -1788,9 +1822,9 @@ template void Index> manager(_query_scratch); auto scratch = manager.scratch_space(); + std::vector pruned_list; if (_filtered_index) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index dc6df02c94..f6f9eb10bc 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -31,7 +31,12 @@ if (NOT Boost_FOUND) message(FATAL_ERROR "Couldn't find Boost dependency") endif() -set(DISKANN_UNIT_TEST_SOURCES main.cpp index_write_parameters_builder_tests.cpp unified_index_tests.cpp filter_match_proxy_tests.cpp) +set(DISKANN_UNIT_TEST_SOURCES + main.cpp + index_write_parameters_builder_tests.cpp + unified_index_tests.cpp + build_optimization_tests.cpp + filter_match_proxy_tests.cpp) # Link the tests against the static DiskANN core (diskann_s) rather than the # DLL, so internal-only symbols don't need to be exported from the DLL just to @@ -97,4 +102,3 @@ foreach(variant ${INSERT_FAST_SIMD_VARIANTS}) add_test(NAME ${v_target} COMMAND ${v_target}) endforeach() - diff --git a/tests/build_optimization_tests.cpp b/tests/build_optimization_tests.cpp new file mode 100644 index 0000000000..8ccf1c0c88 --- /dev/null +++ b/tests/build_optimization_tests.cpp @@ -0,0 +1,327 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ann_exception.h" +#include "in_mem_graph_store.h" +#include "index.h" +#include "parameters.h" +#include "utils.h" + +using namespace diskann; + +namespace +{ +namespace fs = std::filesystem; + +class TestDirectory +{ + public: + TestDirectory() + { + static std::atomic sequence{0}; + const auto timestamp = std::chrono::steady_clock::now().time_since_epoch().count(); + _path = "build_optimization_test_" + std::to_string(timestamp) + "_" + + std::to_string(sequence.fetch_add(1)); + if (!fs::create_directory(_path)) + { + throw std::runtime_error("Failed to create test directory: " + _path.string()); + } + } + + ~TestDirectory() + { + std::error_code error; + fs::remove_all(_path, error); + } + + std::string path(const char *name) const + { + return (_path / fs::path(name)).string(); + } + + private: + fs::path _path; +}; + +std::string read_file(const std::string &path) +{ + std::ifstream input(path, std::ios::binary); + if (!input.is_open()) + { + throw std::runtime_error("Failed to open test file: " + path); + } + return std::string(std::istreambuf_iterator(input), std::istreambuf_iterator()); +} + +std::string line_ending() +{ +#ifdef _WIN32 + return "\r\n"; +#else + return "\n"; +#endif +} + +template ValueT read_value(std::ifstream &input) +{ + ValueT value{}; + input.read(reinterpret_cast(&value), sizeof(value)); + if (!input) + { + throw std::runtime_error("Failed to read graph test fixture"); + } + return value; +} + +void write_float_bin(const std::string &path, uint32_t point_count, uint32_t dimension) +{ + std::ofstream output(path, std::ios::binary | std::ios::trunc); + if (!output.is_open()) + { + throw std::runtime_error("Failed to create data fixture"); + } + + const int32_t points = static_cast(point_count); + const int32_t dimensions = static_cast(dimension); + output.write(reinterpret_cast(&points), sizeof(points)); + output.write(reinterpret_cast(&dimensions), sizeof(dimensions)); + for (uint32_t point = 0; point < point_count; ++point) + { + for (uint32_t coordinate = 0; coordinate < dimension; ++coordinate) + { + uint64_t bits = 42ULL * 1103515245ULL + point * 12345ULL + + coordinate * 7919ULL + 12345ULL; + bits ^= bits >> 21; + bits *= 2685821657736338717ULL; + bits ^= bits >> 31; + const float value = + (static_cast(static_cast(bits)) / 2147483648.0f) - 1.0f; + output.write(reinterpret_cast(&value), sizeof(value)); + } + } +} +} // namespace + +BOOST_AUTO_TEST_CASE(ConvertLabelsPreservesLayoutAndMapping) +{ + TestDirectory directory; + const std::string input_path = directory.path("labels_input.txt"); + const std::string output_path = directory.path("labels_output.txt"); + const std::string map_path = directory.path("labels_map.txt"); + + { + std::ofstream input(input_path, std::ios::binary); + BOOST_REQUIRE(input.is_open()); + input << "red,blue\n" + "blue\n" + "green,red\n"; + } + + uint32_t universal_label_id = 0; + convert_labels_string_to_int(input_path, output_path, map_path, "blue", universal_label_id); + + const std::string expected = + "1,2" + line_ending() + "2" + line_ending() + "3,1" + line_ending(); + BOOST_REQUIRE_EQUAL(read_file(output_path), expected); + BOOST_REQUIRE_EQUAL(universal_label_id, 2); + + std::unordered_map labels; + std::ifstream map_file(map_path); + BOOST_REQUIRE(map_file.is_open()); + for (std::string line; std::getline(map_file, line);) + { + std::istringstream entry(line); + std::string label; + uint32_t id = 0; + BOOST_REQUIRE(static_cast(std::getline(entry, label, '\t'))); + BOOST_REQUIRE(entry >> id); + labels.emplace(std::move(label), id); + } + BOOST_REQUIRE_EQUAL(labels.at("red"), 1); + BOOST_REQUIRE_EQUAL(labels.at("blue"), 2); + BOOST_REQUIRE_EQUAL(labels.at("green"), 3); +} + +BOOST_AUTO_TEST_CASE(ConvertLabelsReportsOpenFailures) +{ + TestDirectory directory; + const std::string missing_path = directory.path("missing_labels.txt"); + const std::string input_path = directory.path("valid_labels.txt"); + const std::string output_path = directory.path("failure_output.txt"); + const std::string map_path = directory.path("failure_map.txt"); + + uint32_t universal_label_id = 0; + BOOST_REQUIRE_THROW( + convert_labels_string_to_int(missing_path, output_path, map_path, "", universal_label_id), + ANNException); + + { + std::ofstream input(input_path); + BOOST_REQUIRE(input.is_open()); + input << "label\n"; + } + + BOOST_REQUIRE_THROW( + convert_labels_string_to_int(input_path, ".", map_path, "", universal_label_id), + ANNException); + BOOST_REQUIRE_THROW( + convert_labels_string_to_int(input_path, output_path, ".", "", universal_label_id), + ANNException); +} + +BOOST_AUTO_TEST_CASE(CopyAndMoveFilePreserveExpectedOwnership) +{ + TestDirectory directory; + const std::string copy_source = directory.path("copy_source.bin"); + const std::string copy_destination = directory.path("copy_destination.bin"); + const std::string move_source = directory.path("move_source.bin"); + const std::string move_destination = directory.path("move_destination.bin"); + + const std::string binary_content("new\0bytes", 9); + { + std::ofstream source(copy_source, std::ios::binary); + source.write(binary_content.data(), binary_content.size()); + std::ofstream destination(copy_destination, std::ios::binary); + destination << "old-longer-content"; + } + + copy_file(copy_source, copy_destination); + BOOST_REQUIRE(fs::exists(copy_source)); + BOOST_REQUIRE_EQUAL(read_file(copy_destination), binary_content); + + { + std::ofstream source(move_source, std::ios::binary); + source << "replacement"; + std::ofstream destination(move_destination, std::ios::binary); + destination << "old-longer-content"; + } + + move_file(move_source, move_destination); + BOOST_REQUIRE(!fs::exists(move_source)); + BOOST_REQUIRE_EQUAL(read_file(move_destination), "replacement"); +} + +BOOST_AUTO_TEST_CASE(InMemoryGraphStorePreservesLegacyBinaryLayout) +{ + TestDirectory directory; + const std::string graph_path = directory.path("graph.bin"); + + InMemGraphStore graph(4, 3); + std::vector neighbors0{1, 2}; + std::vector neighbors1; + std::vector neighbors2{0, 1, 3}; + std::vector neighbors3{2}; + graph.set_neighbours(0, neighbors0); + graph.set_neighbours(1, neighbors1); + graph.set_neighbours(2, neighbors2); + graph.set_neighbours(3, neighbors3); + + constexpr uint64_t expected_size = 64; + BOOST_REQUIRE_EQUAL(graph.store(graph_path, 4, 99, 2), expected_size); + BOOST_REQUIRE_EQUAL(fs::file_size(graph_path), expected_size); + + std::ifstream input(graph_path, std::ios::binary); + BOOST_REQUIRE(input.is_open()); + BOOST_REQUIRE_EQUAL(read_value(input), expected_size); + BOOST_REQUIRE_EQUAL(read_value(input), 3); + BOOST_REQUIRE_EQUAL(read_value(input), 2); + BOOST_REQUIRE_EQUAL(read_value(input), 1); + + const std::vector> expected{ + neighbors0, neighbors1, neighbors2, neighbors3}; + for (const auto &neighbors : expected) + { + BOOST_REQUIRE_EQUAL(read_value(input), neighbors.size()); + for (const uint32_t neighbor : neighbors) + { + BOOST_REQUIRE_EQUAL(read_value(input), neighbor); + } + } + BOOST_REQUIRE_EQUAL(input.peek(), std::char_traits::eof()); +} + +BOOST_AUTO_TEST_CASE(FilteredIndexSavePreservesNumericLabelLayout) +{ + constexpr uint32_t point_count = 64; + constexpr uint32_t dimension = 8; + constexpr uint32_t max_degree = 8; + constexpr uint32_t search_list_size = 16; + + TestDirectory directory; + const std::string data_path = directory.path("index_data.bin"); + const std::string raw_labels_path = directory.path("index_raw_labels.txt"); + const std::string index_prefix = directory.path("filtered_index"); + + write_float_bin(data_path, point_count, dimension); + + std::string expected_labels; + { + std::ofstream labels(raw_labels_path, std::ios::binary); + BOOST_REQUIRE(labels.is_open()); + for (uint32_t point = 0; point < point_count; ++point) + { + if (point == 0) + { + labels << "red,blue\n"; + expected_labels += "1,2" + line_ending(); + } + else if ((point & 1U) == 0) + { + labels << "red\n"; + expected_labels += "1" + line_ending(); + } + else + { + labels << "blue\n"; + expected_labels += "2" + line_ending(); + } + } + } + + auto write_parameters = std::make_shared( + IndexWriteParametersBuilder(search_list_size, max_degree) + .with_alpha(1.2f) + .with_num_threads(1) + .with_filter_list_size(search_list_size) + .build()); + Index index( + Metric::L2, + dimension, + point_count, + write_parameters, + nullptr, + 0, + false, + false, + false, + false, + 0, + false, + true); + + IndexFilterParams filters = IndexFilterParamsBuilder() + .with_label_file(raw_labels_path) + .with_save_path_prefix(index_prefix) + .build(); + index.build(data_path, point_count, filters); + index.save(index_prefix.c_str()); + + BOOST_REQUIRE_EQUAL(read_file(index_prefix + "_labels.txt"), expected_labels); +}