(DNND) Add dnnd_add_knng_edges.cpp

examples/dnnd_bench.cpp

@@ -269,8 +269,8 @@ void usage(std::string_view exe_name, cout_type &cout) {
   cout << "  -g <string> Ground truth file path" << std::endl;
   cout << "  -o <string> Query result file path" << std::endl;
   cout << "  -b <int>    Batch size (default: 1^25)" << std::endl;
-  cout << "  -D <string> kNNG dump prefix" << std::endl;
-  cout << "  -M          Dump index with distance" << std::endl;
+  cout << "  -G <string> kNNG dump prefix" << std::endl;
+  cout << "  -D          Dump index with distance" << std::endl;
   cout << "  -v          Verbose mode" << std::endl;
   cout << "  -h          Show this message" << std::endl;
 }

include/saltatlas/dnnd/detail/nn_index.hpp

@@ -6,6 +6,7 @@
 #pragma once
 
 #include <algorithm>
+#include <filesystem>
 #include <fstream>
 #include <iostream>
 #include <memory>
@@ -194,8 +195,9 @@ class nn_index {
   /// neighbor ids. The second line is the dummy value and distances to each
   /// neighbor. The dummy value is just a placeholder so that each neighbor id
   /// and distance pair is stored in the same column.
-  bool dump(const std::string_view filename, bool dump_distance = false) const {
-    std::ofstream ofs(filename.data());
+  bool dump(const std::filesystem::path &filename,
+            bool                         dump_distance = false) const {
+    std::ofstream ofs(filename);
     if (!ofs) {
       std::cerr << "Failed to open the file: " << filename << std::endl;
       return false;
@@ -225,6 +227,64 @@ class nn_index {
     return true;
   }
 
+  /// \brief Load the index from a dump file.
+  /// \param filename The file name to load the index.
+  /// The file must be formatted as described in the `dump` method.
+  /// \param contains_distance If true, the input file contains distances.
+  /// If false, the input file does not contain distances and distance values
+  /// will be uninitialized.
+  /// \param overwrite If true, the existing index is overwritten.
+  /// \return True if the load is successful.
+  bool load(const std::filesystem::path &filename,
+            const bool contains_distance = false, const bool overwrite = true) {
+    if (overwrite) reset();
+
+    std::ifstream ifs(filename);
+    if (!ifs) {
+      std::cerr << "Failed to open the file: " << filename << std::endl;
+      return false;
+    }
+    std::string line;
+    while (std::getline(ifs, line)) {
+      std::istringstream iss(line);
+      id_type            source;
+      iss >> source;
+      if (iss.fail()) {
+        std::cerr << "Failed to read the source id: " << line << std::endl;
+        return false;
+      }
+      if (m_index.count(source) > 0) {
+        std::cerr << "The source id is duplicated: " << source << std::endl;
+        return false;
+      }
+
+      neighbor_list_type neighbors;
+      id_type            neighbor_id;
+      while (iss >> neighbor_id) {
+        neighbors.emplace_back(neighbor_id, distance_type{});
+      }
+      m_index[source] = std::move(neighbors);
+
+      if (contains_distance) {
+        std::getline(ifs, line);
+        std::istringstream iss(line);
+        distance_type      distance;
+        iss >> distance;  // dummy distance
+        std::size_t index = 0;
+        while (iss >> distance) {
+          m_index[source][index++].distance = distance;
+        }
+        if (index != m_index[source].size()) {
+          std::cerr << "The number of distances does not match the number of "
+                       "neighbors"
+                    << std::endl;
+          return false;
+        }
+      }
+    }
+    return true;
+  }
+
  private:
   point_table_type m_index;
 };

utility/CMakeLists.txt

@@ -13,4 +13,7 @@ if (SALTATLAS_USE_METALL)
 
 	add_executable(dnnd_transfer_pm_datastore dnnd_transfer_pm_datastore.cpp)
 	target_link_libraries(dnnd_transfer_pm_datastore PUBLIC saltatlas)
+
+	add_executable(dnnd_add_knng_edges dnnd_add_knng_edges.cpp)
+	target_link_libraries(dnnd_add_knng_edges PUBLIC saltatlas)
 endif()

utility/dnnd_add_knng_edges.cpp

@@ -0,0 +1,156 @@
+// Copyright 2024 Lawrence Livermore National Security, LLC and other
+// saltatlas Project Developers. See the top-level COPYRIGHT file for details.
+//
+// SPDX-License-Identifier: MIT
+
+// Add additional edges to KNNG output files generated by DNND.
+// Takes a list of edges to add and the original KNNG output files.
+// Outputs a new KNNG output file with the additional edges.
+//
+// NOTE: this program reads all data to the root process to simplify the
+// implementation.
+
+#include <unistd.h>
+
+#include <algorithm>
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <string>
+
+#include <ygm/comm.hpp>
+
+#include <saltatlas/dnnd/data_reader.hpp>
+#include <saltatlas/dnnd/detail/distance.hpp>
+#include <saltatlas/dnnd/detail/neighbor.hpp>
+#include <saltatlas/dnnd/detail/nn_index.hpp>
+#include <saltatlas/dnnd/feature_vector.hpp>
+#include <saltatlas/dnnd/utility.hpp>
+#include <saltatlas/point_store.hpp>
+
+using id_t       = uint32_t;
+using point_t    = saltatlas::feature_vector<float>;
+using distance_t = double;
+
+struct option {
+  std::filesystem::path points_dir;         // i
+  std::string           point_file_format;  // p
+  std::filesystem::path knng_dir;           // g
+  std::filesystem::path edges_dir;          // e
+  std::string           output_file_name;   // o
+  std::string           distance_name;      // f
+};
+
+// use getopt to parse command line options
+bool parse_options(int argc, char** argv, option& opt) {
+  int opt_char;
+  while ((opt_char = ::getopt(argc, argv, "i:p:g:e:o:f:")) != -1) {
+    switch (opt_char) {
+      case 'i':
+        opt.points_dir = std::filesystem::path(optarg);
+        break;
+      case 'p':
+        opt.point_file_format = optarg;
+        break;
+      case 'g':
+        opt.knng_dir = std::filesystem::path(optarg);
+        break;
+      case 'e':
+        opt.edges_dir = std::filesystem::path(optarg);
+        break;
+      case 'o':
+        opt.output_file_name = optarg;
+        break;
+      case 'f':
+        opt.distance_name = optarg;
+        break;
+      default:
+        return false;
+    }
+  }
+  return true;
+}
+
+int main(int argc, char** argv) {
+  ygm::comm comm(&argc, &argv);
+
+  saltatlas::point_store<id_t, point_t> point_store;
+  option                                opt;
+  parse_options(argc, argv, opt);
+
+  comm.cout0() << "\n<<Read Points>>" << std::endl;
+  {
+    const auto paths = saltatlas::utility::find_file_paths(opt.points_dir);
+    saltatlas::read_points<id_t, point_t>(
+        paths, opt.point_file_format, false, [&](const id_t& id) { return 0; },
+        point_store, comm);
+    comm.cout0() << "#of points\t" << point_store.size() << std::endl;
+  }
+
+  if (comm.rank0()) {
+    // Read KNNG
+    std::cout << "\n<<Read KNNG>>" << std::endl;
+    using knng_t = saltatlas::dndetail::nn_index<id_t, distance_t>;
+    knng_t knng;
+    {
+      const auto paths = saltatlas::utility::find_file_paths(opt.knng_dir);
+      for (const auto& path : paths) {
+        std::cout << "Read KNNG file: " << path << std::endl;
+        if (!knng.load(path, true, false)) {
+          std::cerr << "Failed to read KNNG file: " << path << std::endl;
+          return 1;
+        }
+      }
+    }
+
+    // Add edges
+    std::cout << "\n<<Add Edges>>" << std::endl;
+    {
+      auto distance_func =
+          saltatlas::distance::distance_function<point_t, distance_t>(
+              opt.distance_name);
+
+      const auto paths = saltatlas::utility::find_file_paths(opt.edges_dir);
+      for (const auto& path : paths) {
+        std::cout << "Read edge file: " << path << std::endl;
+        std::ifstream ifs(path);
+        if (!ifs.is_open()) {
+          std::cerr << "Cannot open file: " << path << std::endl;
+          return 1;
+        }
+        for (std::string line; std::getline(ifs, line);) {
+          std::istringstream iss(line);
+          id_t               src, dst;
+          iss >> src >> dst;
+
+          if (!point_store.contains(src) || !point_store.contains(dst)) {
+            std::cerr << "Invalid edge: " << src << " " << dst << std::endl;
+            return EXIT_FAILURE;
+          }
+          const auto dist = distance_func(point_store[src], point_store[dst]);
+          knng.insert(src, knng_t::neighbor_type{dst, dist});
+        }
+      }
+
+      // Sort neighbors by distance
+      for (auto pitr = knng.points_begin(), end = knng.points_end();
+           pitr != end; ++pitr) {
+        knng.sort_neighbors(pitr->first);
+      }
+    }
+
+    // Write KNNG
+    std::cout << "\n<<Write KNNG>>" << std::endl;
+    {
+      std::cout << "Write KNNG file: " << opt.output_file_name << std::endl;
+      if (!knng.dump(opt.output_file_name, true)) {
+        std::cerr << "Failed to write KNNG file: " << opt.output_file_name
+                  << std::endl;
+        return 1;
+      }
+    }
+  }
+  comm.cf_barrier();
+
+  return EXIT_SUCCESS;
+}