Add CAGRA build heuristics independent of HNSW parameters

c/include/cuvs/neighbors/cagra.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -299,6 +299,27 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_
                                                 enum cuvsCagraHnswHeuristicType heuristic,
                                                 cuvsDistanceType metric);
 
+/**
+ * @brief Create CAGRA index parameters heuristically tuned for a dataset
+ *
+ * This factory function selects the graph build algorithm and its parameters based on the shape of
+ * the dataset.
+ *
+ * @param[out] params The CAGRA index params to populate
+ * @param[in] n_rows Number of rows in the dataset
+ * @param[in] dim Number of dimensions in the dataset
+ * @param[in] graph_degree Degree of the output graph
+ * @param[in] metric Distance metric to use
+ * @param[in] build_quality Higher values increase build quality (and cost) up to a point
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromDataset(cuvsCagraIndexParams_t params,
+                                                int64_t n_rows,
+                                                int64_t dim,
+                                                size_t graph_degree,
+                                                cuvsDistanceType metric,
+                                                size_t build_quality);
+
 /**
  * @}
  */

c/src/neighbors/cagra.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -834,6 +834,22 @@ extern "C" cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_t
   });
 }
 
+extern "C" cuvsError_t cuvsCagraIndexParamsFromDataset(cuvsCagraIndexParams_t params,
+                                                       int64_t n_rows,
+                                                       int64_t dim,
+                                                       size_t graph_degree,
+                                                       cuvsDistanceType metric,
+                                                       size_t build_quality)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto cpp_metric = static_cast<cuvs::distance::DistanceType>((int)metric);
+    auto cpp_params = cuvs::neighbors::cagra::index_params::from_dataset(
+      raft::matrix_extent<int64_t>(n_rows, dim), graph_degree, cpp_metric, build_quality);
+
+    _populate_cagra_index_params_from_cpp(params, cpp_params);
+  });
+}
+
 extern "C" cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params)
 {
   return cuvs::core::translate_exceptions(

cpp/include/cuvs/neighbors/cagra.hpp

@@ -112,6 +112,12 @@ namespace graph_build_params = cuvs::neighbors::graph_build_params;
  * @{
  */
 
+using graph_build_params_t = std::variant<std::monostate,
+                                          graph_build_params::ivf_pq_params,
+                                          graph_build_params::nn_descent_params,
+                                          graph_build_params::ace_params,
+                                          graph_build_params::iterative_search_params>;
+
 /**
  * @brief A strategy for selecting the graph build parameters based on similar HNSW index
  * parameters.
@@ -181,12 +187,7 @@ struct index_params : cuvs::neighbors::index_params {
    * cagra::graph_build_params::iterative_search_params();
    * @endcode
    */
-  std::variant<std::monostate,
-               graph_build_params::ivf_pq_params,
-               graph_build_params::nn_descent_params,
-               graph_build_params::ace_params,
-               graph_build_params::iterative_search_params>
-    graph_build_params;
+  graph_build_params_t graph_build_params;
   /**
    * Whether to use MST optimization to guarantee graph connectivity.
    */
@@ -223,6 +224,55 @@ struct index_params : cuvs::neighbors::index_params {
    */
   bool attach_dataset_on_build = true;
 
+  /**
+   * @brief Select the graph build algorithm and its parameters for a dataset.
+   *
+   * This is the main CAGRA build heuristic: it chooses between NN-descent and IVF-PQ based on the
+   * dataset size and tunes their parameters based on the target intermediate graph degree and the
+   * requested build quality. It returns the `graph_build_params` variant only; the caller is
+   * responsible for setting `graph_degree` / `intermediate_graph_degree`.
+   *
+   * @param dataset The shape of the input dataset
+   * @param intermediate_graph_degree The intermediate (kNN) graph degree the build should target.
+   *  Note: the intermediate graph degree must be not smaller than the output graph degree; a good
+   *        practice is to have it 1.5x to 2x of the desired graph_degree and a multiple of 32.
+   * @param metric The distance metric to search
+   * @param build_quality Higher values increase the build quality (and cost) up to a point.
+   *        Any value is valid, but values below 20 are the most practical (default = 7).
+   */
+  static graph_build_params_t graph_build_heuristic(
+    raft::matrix_extent<int64_t> dataset,
+    size_t intermediate_graph_degree,
+    cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded,
+    size_t build_quality                = 7);
+
+  /**
+   * @brief Create CAGRA index parameters heuristically tuned for a dataset.
+   *
+   * Returns default CAGRA `index_params` with `graph_build_params` selected by
+   * `graph_build_heuristic` for the given dataset.
+   *
+   * @param dataset The shape of the input dataset
+   * @param graph_degree Degree of the output graph.
+   * @param metric The distance metric to search
+   * @param build_quality Higher values increase the build quality (and cost) up to a point.
+   *        Any value is valid, but values below 20 are the most practical (default = 7).
+   *
+   * Usage example:
+   * @code{.cpp}
+   *   using namespace cuvs::neighbors;
+   *   raft::resources res;
+   *   auto dataset = raft::make_device_matrix<float, int64_t>(res, N, D);
+   *   auto cagra_params = cagra::index_params::from_dataset(dataset.extents());
+   *   auto cagra_index = cagra::build(res, cagra_params, dataset);
+   * @endcode
+   */
+  static cagra::index_params from_dataset(
+    raft::matrix_extent<int64_t> dataset,
+    size_t graph_degree                 = 64,
+    cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded,
+    size_t build_quality                = 7);
+
   /**
    * @brief Create a CAGRA index parameters compatible with HNSW index
    *

cpp/src/neighbors/cagra.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -10,27 +10,40 @@
 
 namespace cuvs::neighbors::cagra {
 
-inline auto graph_params_heuristic(raft::matrix_extent<int64_t> dataset,
-                                   int intermediate_graph_degree,
-                                   int ef_construction,
-                                   cuvs::distance::DistanceType metric)
-  -> decltype(index_params::graph_build_params)
+graph_build_params_t index_params::graph_build_heuristic(raft::matrix_extent<int64_t> dataset,
+                                                         size_t intermediate_graph_degree,
+                                                         cuvs::distance::DistanceType metric,
+                                                         size_t build_quality)
 {
   if (dataset.extent(0) < int64_t(1e6)) {
     // Use NN descent for smaller datasets
     auto nn_descent_params =
       graph_build_params::nn_descent_params(intermediate_graph_degree, metric);
-    nn_descent_params.max_iterations = 5 + ef_construction / 16;
+    nn_descent_params.max_iterations = 5 + build_quality;
     return nn_descent_params;
   } else {
     // Otherwise, use IVF-PQ
     auto ivf_pq_params = cuvs::neighbors::graph_build_params::ivf_pq_params(dataset, metric);
     ivf_pq_params.search_params.n_probes =
-      std::round(2 + std::sqrt(ivf_pq_params.build_params.n_lists) / 20 + ef_construction / 16);
+      std::round(2 + std::sqrt(ivf_pq_params.build_params.n_lists) / 20 + build_quality);
     return ivf_pq_params;
   }
 }
 
+cagra::index_params index_params::from_dataset(raft::matrix_extent<int64_t> dataset,
+                                               size_t graph_degree,
+                                               cuvs::distance::DistanceType metric,
+                                               size_t build_quality)
+{
+  cagra::index_params params;
+  params.metric                    = metric;
+  params.graph_degree              = graph_degree;
+  params.intermediate_graph_degree = graph_degree * 3 / 2;
+  params.graph_build_params =
+    graph_build_heuristic(dataset, params.intermediate_graph_degree, metric, build_quality);
+  return params;
+}
+
 cagra::index_params index_params::from_hnsw_params(raft::matrix_extent<int64_t> dataset,
                                                    int M,
                                                    int ef_construction,
@@ -50,7 +63,7 @@ cagra::index_params index_params::from_hnsw_params(raft::matrix_extent<int64_t>
       break;
   }
   params.graph_build_params =
-    graph_params_heuristic(dataset, params.intermediate_graph_degree, ef_construction, metric);
+    graph_build_heuristic(dataset, params.intermediate_graph_degree, metric, ef_construction / 16);
   return params;
 }
 

cpp/src/neighbors/detail/cagra/cagra_build.cuh

@@ -1383,15 +1383,11 @@ index<T, IdxT> build_ace(raft::resources const& res,
         std::chrono::duration_cast<std::chrono::milliseconds>(read_end - start).count();
 
       // Create index for this partition
-      cuvs::neighbors::cagra::index_params sub_index_params;
-      sub_index_params = cuvs::neighbors::cagra::index_params::from_hnsw_params(
+      auto sub_index_params = cuvs::neighbors::cagra::index_params::from_dataset(
         raft::make_extents<int64_t>(sub_dataset_size, dataset_dim),
-        graph_degree / 2,
-        ef_construction,
-        cuvs::neighbors::cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT,
-        params.metric);
-      // avoid rounding down graph_degree via heuristics doing `(graph_degree / 2) * 2`
-      sub_index_params.graph_degree            = graph_degree;
+        graph_degree,
+        params.metric,
+        ef_construction / 16);
       sub_index_params.attach_dataset_on_build = false;
       sub_index_params.guarantee_connectivity  = params.guarantee_connectivity;
 

java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndexParams.java

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 package com.nvidia.cuvs;
@@ -359,6 +359,12 @@ public static CagraIndexParams fromHnswParams(
         .cagraIndexParamsFromHnswParams(rows, dim, M, efConstruction, heuristic, metric);
   }
 
+  public static CagraIndexParams fromDataset(
+      long rows, long dim, long graphDegree, CuvsDistanceType metric, long buildQuality) {
+    return CuVSProvider.provider()
+        .cagraIndexParamsFromDataset(rows, dim, graphDegree, metric, buildQuality);
+  }
+
   /**
    * Gets the degree of input graph for pruning.
    *

java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java

@@ -246,4 +246,21 @@ CagraIndexParams cagraIndexParamsFromHnswParams(
       int efConstruction,
       CagraIndexParams.HnswHeuristicType heuristic,
       CagraIndexParams.CuvsDistanceType metric);
+
+  /**
+   * Create CAGRA index parameters heuristically tuned for a dataset.
+   *
+   * @param rows The number of rows in the input dataset
+   * @param dim The number of dimensions in the input dataset
+   * @param graphDegree Degree of the output graph
+   * @param metric The distance metric to search
+   * @param buildQuality Higher values increase build quality (and cost) up to a point
+   * @return A new CAGRA index parameters object
+   */
+  CagraIndexParams cagraIndexParamsFromDataset(
+      long rows,
+      long dim,
+      long graphDegree,
+      CagraIndexParams.CuvsDistanceType metric,
+      long buildQuality);
 }

java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java

@@ -105,6 +105,16 @@ public CagraIndexParams cagraIndexParamsFromHnswParams(
     throw new UnsupportedOperationException(reasons);
   }
 
+  @Override
+  public CagraIndexParams cagraIndexParamsFromDataset(
+      long rows,
+      long dim,
+      long graphDegree,
+      CagraIndexParams.CuvsDistanceType metric,
+      long buildQuality) {
+    throw new UnsupportedOperationException(reasons);
+  }
+
   @Override
   public void setLogLevel(Level level) {
     throw new UnsupportedOperationException(reasons);

java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java

@@ -345,6 +345,40 @@ public CagraIndexParams cagraIndexParamsFromHnswParams(
     }
   }
 
+  @Override
+  public CagraIndexParams cagraIndexParamsFromDataset(
+      long rows,
+      long dim,
+      long graphDegree,
+      CagraIndexParams.CuvsDistanceType metric,
+      long buildQuality) {
+    try (var nativeCagraIndexParams = createCagraIndexParams();
+        var ivfPqIndexParams = createIvfPqIndexParams();
+        var ivfPqSearchParams = createIvfPqSearchParams()) {
+
+      // This is already allocated by cuvsCagraIndexParamsCreate,
+      // we just need to populate it.
+      MemorySegment cuvsIvfPqParamsMemorySegment =
+          cuvsCagraIndexParams.graph_build_params(nativeCagraIndexParams.handle());
+      cuvsIvfPqParams.ivf_pq_build_params(cuvsIvfPqParamsMemorySegment, ivfPqIndexParams.handle());
+      cuvsIvfPqParams.ivf_pq_search_params(
+          cuvsIvfPqParamsMemorySegment, ivfPqSearchParams.handle());
+
+      cuvsCagraIndexParams.graph_build_params(
+          nativeCagraIndexParams.handle(), cuvsIvfPqParamsMemorySegment);
+      checkCuVSError(
+          cuvsCagraIndexParamsFromDataset(
+              nativeCagraIndexParams.handle(), rows, dim, graphDegree, metric.value, buildQuality),
+          "cuvsCagraIndexParamsFromDataset");
+
+      return populateCagraIndexParamsFromNative(
+          nativeCagraIndexParams,
+          ivfPqIndexParams,
+          ivfPqSearchParams,
+          cuvsIvfPqParamsMemorySegment);
+    }
+  }
+
   @Override
   public void setLogLevel(Level level) {
     if (level.equals(Level.ALL) || level.equals(Level.FINEST)) {

rust/cuvs-sys/src/bindings.rs

@@ -1323,6 +1323,18 @@ unsafe extern "C" {
         metric: cuvsDistanceType,
     ) -> cuvsError_t;
 }
+unsafe extern "C" {
+    #[must_use]
+    #[doc = " @brief Create CAGRA index parameters heuristically tuned for a dataset\n\n This factory function selects the graph build algorithm and its parameters based on the shape of\n the dataset.\n\n @param[out] params The CAGRA index params to populate\n @param[in] n_rows Number of rows in the dataset\n @param[in] dim Number of dimensions in the dataset\n @param[in] graph_degree Degree of the output graph\n @param[in] metric Distance metric to use\n @param[in] build_quality Higher values increase build quality (and cost) up to a point\n @return cuvsError_t"]
+    pub fn cuvsCagraIndexParamsFromDataset(
+        params: cuvsCagraIndexParams_t,
+        n_rows: i64,
+        dim: i64,
+        graph_degree: usize,
+        metric: cuvsDistanceType,
+        build_quality: usize,
+    ) -> cuvsError_t;
+}
 #[doc = " @defgroup cagra_c_extend_params C API for CUDA ANN Graph-based nearest neighbor search\n @{\n/\n/**\n @brief Supplemental parameters to extend CAGRA Index\n"]
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]