Add Augmented Core Extraction Algorithm

c/include/cuvs/neighbors/cagra.h

@@ -34,7 +34,9 @@ enum cuvsCagraGraphBuildAlgo {
   /* Experimental, use NN-Descent to build all-neighbors knn graph */
   NN_DESCENT = 2,
   /* Experimental, use iterative cagra search and optimize to build the knn graph */
-  ITERATIVE_CAGRA_SEARCH = 3
+  ITERATIVE_CAGRA_SEARCH = 3,
+  /* Use ACE (Augmented Core Extraction) to build the graph */
+  ACE = 4
 };
 
 /** Parameters for VPQ compression. */
@@ -84,6 +86,36 @@ struct cuvsIvfPqParams {
 
 typedef struct cuvsIvfPqParams* cuvsIvfPqParams_t;
 
+/** Parameters for ACE (Augmented Core Extraction) graph build */
+struct cuvsAceParams {
+  /**
+   * Number of partitions for ACE (Augmented Core Extraction) partitioned build.
+   * Small values might improve recall but potentially degrade performance and
+   * increase memory usage. Partitions should not be too small to prevent issues
+   * in KNN graph construction. 100k - 5M vectors per partition is recommended
+   * depending on the available host and GPU memory.
+   */
+  size_t npartitions;
+  /**
+   * The index quality for the ACE build.
+   * Bigger values increase the index quality. At some point, increasing this will no longer
+   * improve the quality.
+   */
+  size_t ef_construction;
+  /**
+   * Directory to store ACE build artifacts (e.g., KNN graph, optimized graph).
+   * Used when `npartitions` > 1 or `use_disk` is true.
+   */
+  const char* build_dir;
+  /**
+   * Whether to use disk-based storage for ACE build.
+   * When true, enables disk-based operations for memory-efficient graph construction.
+   */
+  bool use_disk;
+};
+
+typedef struct cuvsAceParams* cuvsAceParams_t;
+
 /**
  * @brief Supplemental parameters to build CAGRA Index
  *
@@ -106,9 +138,12 @@ struct cuvsCagraIndexParams {
    */
   cuvsCagraCompressionParams_t compression;
   /**
-   * Optional: specify ivf pq params when `build_algo = IVF_PQ`
+   * Optional: specify graph build params based on build_algo
+   * - IVF_PQ: cuvsIvfPqParams_t
+   * - ACE: cuvsAceParams_t
+   * - Others: nullptr
    */
-  cuvsIvfPqParams_t graph_build_params;
+  void* graph_build_params;
 };
 
 typedef struct cuvsCagraIndexParams* cuvsCagraIndexParams_t;
@@ -145,6 +180,22 @@ cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t* param
  */
 cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params);
 
+/**
+ * @brief Allocate ACE params, and populate with default values
+ *
+ * @param[in] params cuvsAceParams_t to allocate
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsAceParamsCreate(cuvsAceParams_t* params);
+
+/**
+ * @brief De-allocate ACE params
+ *
+ * @param[in] params
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsAceParamsDestroy(cuvsAceParams_t params);
+
 /**
  * @}
  */
@@ -365,6 +416,25 @@ cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size);
  */
 cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree);
 
+/**
+ * @brief Check if the CAGRA index is stored on disk
+ *
+ * @param[in] index CAGRA index
+ * @param[out] on_disk return true if index is on disk, false otherwise
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraIndexIsOnDisk(cuvsCagraIndex_t index, bool* on_disk);
+
+/**
+ * @brief Get the file directory where the CAGRA index is stored (if on disk)
+ *
+ * @param[in] index CAGRA index
+ * @param[out] file_directory return file directory path (caller must free)
+ * @param[out] length length of the file_directory string
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraIndexGetFileDirectory(cuvsCagraIndex_t index, char** file_directory, size_t* length);
+
 /**
  * @brief Returns a view of the CAGRA dataset
  *

c/src/neighbors/cagra.cpp

@@ -4,6 +4,7 @@
  */
 
 #include <cstdint>
+#include <cstring>
 #include <dlpack/dlpack.h>
 
 #include <raft/core/error.hpp>
@@ -15,6 +16,7 @@
 #include <cuvs/neighbors/cagra.h>
 #include <cuvs/neighbors/common.h>
 #include <cuvs/neighbors/cagra.hpp>
+#include <cuvs/neighbors/graph_build_types.hpp>
 
 #include "../core/exceptions.hpp"
 #include "../core/interop.hpp"
@@ -28,6 +30,7 @@ static void _set_graph_build_params(
   std::variant<std::monostate,
                cuvs::neighbors::cagra::graph_build_params::ivf_pq_params,
                cuvs::neighbors::cagra::graph_build_params::nn_descent_params,
+               cuvs::neighbors::cagra::graph_build_params::ace_params,
                cuvs::neighbors::cagra::graph_build_params::iterative_search_params>& out_params,
   cuvsCagraIndexParams& params,
   cuvsCagraGraphBuildAlgo algo,
@@ -79,6 +82,18 @@ static void _set_graph_build_params(
       out_params               = nn_params;
       break;
     }
+    case cuvsCagraGraphBuildAlgo::ACE: {
+      cuvs::neighbors::cagra::graph_build_params::ace_params ace_p;
+      if (params.graph_build_params) {
+        auto ace_params_c             = static_cast<cuvsAceParams*>(params.graph_build_params);
+        ace_p.npartitions     = ace_params_c->npartitions;
+        ace_p.ef_construction = ace_params_c->ef_construction;
+        ace_p.build_dir       = std::string(ace_params_c->build_dir);
+        ace_p.use_disk        = ace_params_c->use_disk;
+      }
+      out_params = ace_p;
+      break;
+    }
     case cuvsCagraGraphBuildAlgo::ITERATIVE_CAGRA_SEARCH: {
       cuvs::neighbors::cagra::graph_build_params::iterative_search_params p;
       out_params = p;
@@ -432,6 +447,28 @@ extern "C" cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int6
   });
 }
 
+extern "C" cuvsError_t cuvsCagraIndexIsOnDisk(cuvsCagraIndex_t index, bool* on_disk)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<float, uint32_t>*>(index->addr);
+    *on_disk       = index_ptr->on_disk();
+  });
+}
+
+extern "C" cuvsError_t cuvsCagraIndexGetFileDirectory(cuvsCagraIndex_t index,
+                                                      char** file_directory,
+                                                      size_t* length)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto index_ptr     = reinterpret_cast<cuvs::neighbors::cagra::index<float, uint32_t>*>(index->addr);
+    const auto& dir    = index_ptr->file_directory();
+    *length            = dir.length();
+    *file_directory    = static_cast<char*>(malloc((*length + 1) * sizeof(char)));
+    std::strncpy(*file_directory, dir.c_str(), *length);
+    (*file_directory)[*length] = '\0';
+  });
+}
+
 extern "C" cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, DLManagedTensor* dataset)
 {
   return cuvs::core::translate_exceptions([=] {
@@ -641,7 +678,26 @@ extern "C" cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params
 extern "C" cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params)
 {
   return cuvs::core::translate_exceptions([=] {
-    delete params->graph_build_params;
+    // Delete graph_build_params based on the build algorithm type
+    if (params->graph_build_params != nullptr) {
+      switch (params->build_algo) {
+      case cuvsCagraGraphBuildAlgo::IVF_PQ:
+        delete static_cast<cuvsIvfPqParams *>(params->graph_build_params);
+        break;
+      case cuvsCagraGraphBuildAlgo::ACE: {
+        auto ace_params = static_cast<cuvsAceParams *>(params->graph_build_params);
+        // Free the allocated build directory string
+        if (ace_params->build_dir) { free(const_cast<char*>(ace_params->build_dir)); }
+        delete ace_params;
+        break;
+      }
+      case cuvsCagraGraphBuildAlgo::AUTO_SELECT:
+      case cuvsCagraGraphBuildAlgo::NN_DESCENT:
+      case cuvsCagraGraphBuildAlgo::ITERATIVE_CAGRA_SEARCH:
+        // These algorithms don't have separate parameter structs
+        break;
+      }
+    }
     delete params;
   });
 }
@@ -665,6 +721,32 @@ extern "C" cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionPar
   return cuvs::core::translate_exceptions([=] { delete params; });
 }
 
+extern "C" cuvsError_t cuvsAceParamsCreate(cuvsAceParams_t* params)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto ps = cuvs::neighbors::cagra::graph_build_params::ace_params();
+
+    // Allocate and copy the build directory string
+    const char* build_dir = strdup(ps.build_dir.c_str());
+
+    *params = new cuvsAceParams{.npartitions     = ps.npartitions,
+                                .ef_construction = ps.ef_construction,
+                                .build_dir       = build_dir,
+                                .use_disk        = ps.use_disk};
+  });
+}
+
+extern "C" cuvsError_t cuvsAceParamsDestroy(cuvsAceParams_t params)
+{
+  return cuvs::core::translate_exceptions([=] {
+    if (params) {
+      // Free the allocated build directory string
+      if (params->build_dir) { free(const_cast<char*>(params->build_dir)); }
+      delete params;
+    }
+  });
+}
+
 extern "C" cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params)
 {
   return cuvs::core::translate_exceptions(