LLNL · KIwabuchi · Feb 6, 2025 · Jan 30, 2025 · Jan 30, 2025 · Jan 31, 2025
diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -43,6 +43,9 @@ add_saltatlas_dnnd_example_feature_type(dnnd_bench double)
 
 if (SALTATLAS_USE_METALL)
     add_saltatlas_example(dnnd_advanced)
+    add_saltatlas_example(dnnd_advance_index_build)
+    add_saltatlas_example(dnnd_advance_get_neighbors)
+    add_saltatlas_example(dnnd_advance_query)
 endif ()
 
 add_subdirectory(datasets)
diff --git a/examples/dnnd_advance_get_neighbors.cpp b/examples/dnnd_advance_get_neighbors.cpp
@@ -0,0 +1,146 @@
+// Copyright 2024 Lawrence Livermore National Security, LLC and other
+// saltatlas Project Developers. See the top-level COPYRIGHT file for details.
+//
+// SPDX-License-Identifier: MIT
+
+/// \brief A simple example of building k-NN index (KNN graph) in Metall
+/// datastore. Usage:
+///     cd build
+///     mpirun -n 2 ./example/dnnd_advanced_index_build -p /path/to/points -f l2
+
+#include <filesystem>
+#include <iostream>
+#include <string>
+#include <vector>
+
+#include <ygm/comm.hpp>
+
+#include <saltatlas/dnnd/dnnd_advanced.hpp>
+
+// Point ID type
+using id_t   = uint32_t;
+using dist_t = double;
+
+// Point Type
+using point_type = saltatlas::pm_feature_vector<float>;
+
+struct option_t {
+  std::filesystem::path datastore_path;
+  std::vector<id_t>     point_ids;
+};
+
+bool parse_options(int argc, char **argv, option_t &opt, bool &help) {
+  opt.datastore_path.clear();
+  help = false;
+
+  int n;
+  while ((n = ::getopt(argc, argv, "d:p:h")) != -1) {
+    switch (n) {
+      case 'd':
+        opt.datastore_path = optarg;
+        break;
+
+      case 'p':  // comma separated list of point ids
+      {
+        std::string        point_ids_str = optarg;
+        std::istringstream ss(point_ids_str);
+        std::string        token;
+        while (std::getline(ss, token, ',')) {
+          opt.point_ids.push_back(std::stoi(token));
+        }
+        break;
+      }
+
+      case 'h':
+        help = true;
+        return true;
+
+      default:
+        return false;
+    }
+  }
+
+  if (opt.datastore_path.empty()) {
+    return false;
+  }
+
+  return true;
+}
+
+template <typename cout_type>
+void show_help(cout_type &cout) {
+  cout
+      << "Usage: ./dnnd_advance_show_neighbors [options]\n"
+         "Options:\n"
+         "  -d <string>       The Metall datastore path\n"
+         "  -p <string>       Comma separated list of point IDs (e.g., 0,2,5)\n"
+         "  -h                Show this help message\n";
+}
+
+int main(int argc, char **argv) {
+  ygm::comm comm(&argc, &argv);
+
+  option_t opt;
+  bool     help{false};
+  if (!parse_options(argc, argv, opt, help)) {
+    comm.cerr0() << "Invalid option" << std::endl;
+    show_help(comm.cerr0());
+    return EXIT_FAILURE;
+  }
+  if (help) {
+    show_help(comm.cout0());
+    return 0;
+  }
+  if (!comm.rank0()) {
+    // Only rank 0 request neighbors in this demo
+    opt.point_ids.clear();
+  }
+
+  {
+    saltatlas::dnnd<id_t, point_type, dist_t> g(saltatlas::open_read_only,
+                                                opt.datastore_path, comm);
+
+    // Use the first index for demo
+    const auto index_id = g.get_index_ids().front();
+    comm.cout0() << "Index ID: " << index_id << std::endl;
+
+    {
+      const auto ret =
+          g.get_neighbors(index_id, opt.point_ids.begin(), opt.point_ids.end());
+      for (const auto &item : ret) {
+        const auto &pid       = item.first;
+        const auto &neighbors = item.second;
+        comm.cout0() << "Source point ID: " << pid << std::endl;
+        for (const auto &n : neighbors) {
+          comm.cout0() << n << std::endl;
+        }
+      }
+      comm.cout0() << std::endl;
+    }
+
+    {
+      std::vector<id_t> ids;
+      if (comm.rank() == 0) {
+        ids.push_back(0);
+        ids.push_back(1);
+      }
+      const auto neighbors_and_features =
+          g.get_neighbors_with_features(index_id, ids.begin(), ids.end());
+      for (const auto &[id, item] : neighbors_and_features) {
+        comm.cout0() << "Source point ID: " << id << std::endl;
+        const auto &neighbors = item.first;
+        const auto &features  = item.second;
+        for (int i = 0; i < neighbors.size(); ++i) {
+          comm.cout0() << neighbors[i]
+                       << ", feature = " << saltatlas::to_string(features[i])
+                       << std::endl;
+        }
+        comm.cout0() << std::endl;
+      }
+    }
+
+    comm.cf_barrier();
+  }
+
+  return 0;
+}
diff --git a/examples/dnnd_advance_index_build.cpp b/examples/dnnd_advance_index_build.cpp
@@ -0,0 +1,135 @@
+// Copyright 2024 Lawrence Livermore National Security, LLC and other
+// saltatlas Project Developers. See the top-level COPYRIGHT file for details.
+//
+// SPDX-License-Identifier: MIT
+
+/// \brief A simple example of building k-NN index (KNN graph) in Metall
+/// datastore. Usage:
+///     cd build
+///     mpirun -n 2 ./example/dnnd_advanced_index_build -p /path/to/points -f l2
+
+#include <iostream>
+#include <vector>
+#include <string>
+#include <filesystem>
+
+#include <ygm/comm.hpp>
+
+#include <saltatlas/dnnd/dnnd_advanced.hpp>
+
+// Point ID type
+using id_t   = uint32_t;
+using dist_t = double;
+
+// Point Type
+using point_type = saltatlas::pm_feature_vector<float>;
+
+struct option_t {
+  int                                index_k;
+  std::string                        distance_name;
+  std::vector<std::filesystem::path> point_file_names;
+  std::string                        point_file_format;
+  std::filesystem::path              datastore_path;
+};
+
+bool parse_options(int argc, char **argv, option_t &opt, bool &help) {
+  opt.index_k = 0;
+  opt.distance_name.clear();
+  opt.point_file_names.clear();
+  opt.point_file_format.clear();
+  opt.datastore_path.clear();
+  help = false;
+
+  int n;
+  while ((n = ::getopt(argc, argv, "k:f:p:d:h")) != -1) {
+    switch (n) {
+      case 'k':
+        opt.index_k = std::stoi(optarg);
+        break;
+
+      case 'f':
+        opt.distance_name = optarg;
+        break;
+
+      case 'p':
+        opt.point_file_format = optarg;
+        break;
+
+      case 'd':
+        opt.datastore_path = optarg;
+        break;
+
+      case 'h':
+        help = true;
+        return true;
+
+      default:
+        return false;
+    }
+  }
+
+  for (int index = optind; index < argc; index++) {
+    opt.point_file_names.emplace_back(argv[index]);
+  }
+
+  if (opt.index_k <= 0) {
+    return false;
+  }
+
+  if (opt.distance_name.empty() || opt.point_file_format.empty() ||
+      opt.point_file_names.empty() || opt.datastore_path.empty()) {
+    return false;
+  }
+
+  return true;
+}
+
+template <typename cout_type>
+void show_help(cout_type &cout) {
+  cout << "Usage: ./dnnd_advanced_index_build [options] point_files...\n"
+          "Options:\n"
+          "  -k <int>          The number of neighbors to build the index\n"
+          "  -f <string>       The distance function name\n"
+          "  -p <string>       The point file format\n"
+          "  -d <string>       The Metall datastore path\n"
+          "  -h                Show this help message\n";
+}
+
+int main(int argc, char **argv) {
+  ygm::comm comm(&argc, &argv);
+
+  option_t opt;
+  bool     help{false};
+  if (!parse_options(argc, argv, opt, help)) {
+    comm.cerr0() << "Invalid option" << std::endl;
+    show_help(comm.cerr0());
+    return EXIT_FAILURE;
+  }
+  if (help) {
+    show_help(comm.cout0());
+    return 0;
+  }
+
+  std::error_code ec;
+  std::filesystem::remove_all(opt.datastore_path, ec);
+  comm.cf_barrier();
+  {
+    saltatlas::dnnd<id_t, point_type, dist_t> g(saltatlas::create_only,
+                                                opt.datastore_path, comm);
+
+    g.load_points(opt.point_file_names.begin(), opt.point_file_names.end(),
+                  opt.point_file_format);
+
+    const auto distance_func =
+        saltatlas::distance::distance_function<point_type, dist_t>(
+            opt.distance_name);
+
+    comm.cout0() << "Building index" << std::endl;
+    const auto index_id = g.build(distance_func, opt.index_k);
+
+    comm.cout0() << "Optimizing index" << std::endl;
+    g.optimize(index_id, distance_func);
+  }
+
+  return 0;
+}