[REVIEW] New Dataset API Clarifying Ownership

c/include/cuvs/neighbors/cagra.h

@@ -211,12 +211,6 @@ struct cuvsCagraIndexParams {
   enum cuvsCagraGraphBuildAlgo build_algo;
   /** Number of Iterations to run if building with NN_DESCENT */
   size_t nn_descent_niter;
-  /**
-   * Optional: specify compression parameters if compression is desired.
-   *
-   * NOTE: this is experimental new API, consider it unsafe.
-   */
-  cuvsCagraCompressionParams_t compression;
   /**
    * Optional: specify graph build params based on build_algo
    * - IVF_PQ: cuvsIvfPqParams_t
@@ -466,13 +460,15 @@ CUVS_EXPORT cuvsError_t cuvsCagraSearchParamsDestroy(cuvsCagraSearchParams_t par
  */
 
 /**
- * @brief Struct to hold address of cuvs::neighbors::cagra::index and its active trained dtype
+ * @brief Struct holding the CAGRA index storage address and vector element dtype (DLPack-style)
  *
+ * Matches the usual cuVS C index pattern (`addr` + `dtype`). \p addr points at implementation-owned
+ * storage (not always a bare `cagra::index*`); free only via \ref cuvsCagraIndexDestroy. \p dtype
+ * describes index vector elements for queries and template dispatch.
  */
 typedef struct {
   uintptr_t addr;
   DLDataType dtype;
-
 } cuvsCagraIndex;
 
 typedef cuvsCagraIndex* cuvsCagraIndex_t;

c/src/neighbors/cagra.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #include <cuvs/neighbors/cagra.h>
@@ -15,4 +15,7 @@ void convert_c_index_params(cuvsCagraIndexParams params,
 /// Converts C search params to C++
 void convert_c_search_params(cuvsCagraSearchParams params,
                              cuvs::neighbors::cagra::search_params* out);
+
+/** Resolves `cuvsCagraIndex::addr` to `cagra::index*`; nullptr if the handle is empty. */
+void* cagra_c_api_index_ptr(cuvsCagraIndex const* idx);
 }  // namespace cuvs::neighbors::cagra

c/src/neighbors/hnsw.cpp

@@ -20,6 +20,7 @@
 
 #include "../core/exceptions.hpp"
 #include "../core/interop.hpp"
+#include "cagra.hpp"
 
 namespace {
 
@@ -63,7 +64,8 @@ void _from_cagra(cuvsResources_t res,
                  std::optional<DLManagedTensor*> dataset_tensor)
 {
   auto res_ptr = reinterpret_cast<raft::resources*>(res);
-  auto index   = reinterpret_cast<cuvs::neighbors::cagra::index<T, uint32_t>*>(cagra_index->addr);
+  auto index   = reinterpret_cast<cuvs::neighbors::cagra::device_padded_index<T, uint32_t>*>(
+    cuvs::neighbors::cagra::cagra_c_api_index_ptr(cagra_index));
   auto cpp_params            = cuvs::neighbors::hnsw::index_params();
   cpp_params.hierarchy       = static_cast<cuvs::neighbors::hnsw::HnswHierarchy>(params->hierarchy);
   cpp_params.ef_construction = params->ef_construction;

c/src/neighbors/mg_cagra.cpp

@@ -84,26 +84,26 @@ extern "C" cuvsError_t cuvsMultiGpuCagraIndexDestroy(cuvsMultiGpuCagraIndex_t in
       // Properly clean up the templated inner object based on dtype, like single GPU API
       if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
         auto mg_index_ptr =
-          reinterpret_cast<cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<float, uint32_t>,
+          reinterpret_cast<cuvs::neighbors::mg_index<cuvs::neighbors::cagra::device_padded_index<float, uint32_t>,
                                                      float,
                                                      uint32_t>*>(index->addr);
         delete mg_index_ptr;
       } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
         auto mg_index_ptr =
-          reinterpret_cast<cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<half, uint32_t>,
+          reinterpret_cast<cuvs::neighbors::mg_index<cuvs::neighbors::cagra::device_padded_index<half, uint32_t>,
                                                      half,
                                                      uint32_t>*>(index->addr);
         delete mg_index_ptr;
       } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
         auto mg_index_ptr = reinterpret_cast<
           cuvs::neighbors::
-            mg_index<cuvs::neighbors::cagra::index<int8_t, uint32_t>, int8_t, uint32_t>*>(
+            mg_index<cuvs::neighbors::cagra::device_padded_index<int8_t, uint32_t>, int8_t, uint32_t>*>(
           index->addr);
         delete mg_index_ptr;
       } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
         auto mg_index_ptr = reinterpret_cast<
           cuvs::neighbors::
-            mg_index<cuvs::neighbors::cagra::index<uint8_t, uint32_t>, uint8_t, uint32_t>*>(
+            mg_index<cuvs::neighbors::cagra::device_padded_index<uint8_t, uint32_t>, uint8_t, uint32_t>*>(
           index->addr);
         delete mg_index_ptr;
       }
@@ -159,7 +159,7 @@ void* _mg_build(cuvsResources_t res,
   auto mds          = cuvs::core::from_dlpack<mdspan_type>(dataset_tensor);
 
   auto mg_index =
-    new cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<T, uint32_t>, T, uint32_t>(
+    new cuvs::neighbors::mg_index<cuvs::neighbors::cagra::device_padded_index<T, uint32_t>, T, uint32_t>(
       cuvs::neighbors::cagra::build(*res_ptr, mg_params, mds));
 
   return mg_index;
@@ -175,7 +175,7 @@ void _mg_search(cuvsResources_t res,
 {
   auto res_ptr      = reinterpret_cast<raft::resources*>(res);
   auto mg_index_ptr = reinterpret_cast<
-    cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<T, uint32_t>, T, uint32_t>*>(
+    cuvs::neighbors::mg_index<cuvs::neighbors::cagra::device_padded_index<T, uint32_t>, T, uint32_t>*>(
     index.addr);
 
   auto mg_search_params =
@@ -202,7 +202,7 @@ void _mg_extend(cuvsResources_t res,
 {
   auto res_ptr      = reinterpret_cast<raft::resources*>(res);
   auto mg_index_ptr = reinterpret_cast<
-    cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<T, uint32_t>, T, uint32_t>*>(
+    cuvs::neighbors::mg_index<cuvs::neighbors::cagra::device_padded_index<T, uint32_t>, T, uint32_t>*>(
     index.addr);
 
   using vectors_mdspan_type = raft::host_matrix_view<const T, int64_t, raft::row_major>;
@@ -222,7 +222,7 @@ void _mg_serialize(cuvsResources_t res, cuvsMultiGpuCagraIndex index, const char
 {
   auto res_ptr      = reinterpret_cast<raft::resources*>(res);
   auto mg_index_ptr = reinterpret_cast<
-    cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<T, uint32_t>, T, uint32_t>*>(
+    cuvs::neighbors::mg_index<cuvs::neighbors::cagra::device_padded_index<T, uint32_t>, T, uint32_t>*>(
     index.addr);
 
   cuvs::neighbors::cagra::serialize(*res_ptr, *mg_index_ptr, std::string(filename));
@@ -233,7 +233,7 @@ void* _mg_deserialize(cuvsResources_t res, const char* filename)
 {
   auto res_ptr = reinterpret_cast<raft::resources*>(res);
   auto mg_index =
-    new cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<T, uint32_t>, T, uint32_t>(
+    new cuvs::neighbors::mg_index<cuvs::neighbors::cagra::device_padded_index<T, uint32_t>, T, uint32_t>(
       cuvs::neighbors::cagra::deserialize<T, uint32_t>(*res_ptr, std::string(filename)));
 
   return mg_index;
@@ -244,7 +244,7 @@ void* _mg_distribute(cuvsResources_t res, const char* filename)
 {
   auto res_ptr = reinterpret_cast<raft::resources*>(res);
   auto mg_index =
-    new cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<T, uint32_t>, T, uint32_t>(
+    new cuvs::neighbors::mg_index<cuvs::neighbors::cagra::device_padded_index<T, uint32_t>, T, uint32_t>(
       cuvs::neighbors::cagra::distribute<T, uint32_t>(*res_ptr, std::string(filename)));
 
   return mg_index;

c/src/neighbors/tiered_index.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -71,7 +71,7 @@ void* _build(cuvsResources_t res, cuvsTieredIndexParams params, DLManagedTensor*
     case CUVS_TIERED_INDEX_ALGO_CAGRA: {
       auto build_params = tiered_index::index_params<cagra::index_params>();
       convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params);
-      return new tiered_index::index<cagra::index<T, uint32_t>>(
+      return new tiered_index::index<cagra::device_padded_index<T, uint32_t>>(
         tiered_index::build(*res_ptr, build_params, mds));
     }
     case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: {
@@ -219,7 +219,7 @@ extern "C" cuvsError_t cuvsTieredIndexDestroy(cuvsTieredIndex_t index_c_ptr)
       switch (index.algo) {
         case CUVS_TIERED_INDEX_ALGO_CAGRA: {
           auto index_ptr =
-            reinterpret_cast<tiered_index::index<cagra::index<float, uint32_t>>*>(index.addr);
+            reinterpret_cast<tiered_index::index<cagra::device_padded_index<float, uint32_t>>*>(index.addr);
           delete index_ptr;
           break;
         }
@@ -292,7 +292,7 @@ extern "C" cuvsError_t cuvsTieredIndexSearch(cuvsResources_t res,
 
     switch (index.algo) {
       case CUVS_TIERED_INDEX_ALGO_CAGRA: {
-        _search<cagra::index<float, uint32_t>>(
+        _search<cagra::device_padded_index<float, uint32_t>>(
           res, search_params, index, queries_tensor, neighbors_tensor, distances_tensor, filter);
         break;
       }
@@ -336,7 +336,7 @@ extern "C" cuvsError_t cuvsTieredIndexExtend(cuvsResources_t res,
     auto index = *index_c_ptr;
     switch (index.algo) {
       case CUVS_TIERED_INDEX_ALGO_CAGRA: {
-        _extend<cagra::index<float, uint32_t>>(res, new_vectors, index);
+        _extend<cagra::device_padded_index<float, uint32_t>>(res, new_vectors, index);
         break;
       }
       case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: {
@@ -363,7 +363,7 @@ extern "C" cuvsError_t cuvsTieredIndexMerge(cuvsResources_t res,
 
     switch (indices[0]->algo) {
       case CUVS_TIERED_INDEX_ALGO_CAGRA: {
-        _merge<cagra::index<float, uint32_t>>(res, *params, indices, num_indices, output_index);
+        _merge<cagra::device_padded_index<float, uint32_t>>(res, *params, indices, num_indices, output_index);
         break;
       }
       case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: {

c/tests/neighbors/ann_cagra_c.cu

@@ -165,12 +165,23 @@ TEST(CagraC, BuildExtendSearch)
     (main_data_size + additional_data_size + num_queries) * dimensions, stream);
   rmm::device_uvector<int32_t> random_labels_d(
     (main_data_size + additional_data_size + num_queries) * dimensions, stream);
-  raft::random::make_blobs(random_data_d.data(),
-                           random_labels_d.data(),
-                           main_data_size + additional_data_size + num_queries,
-                           dimensions,
-                           10,
-                           stream);
+
+  raft::random::make_blobs<float, int32_t>(
+    random_data_d.data(),
+    random_labels_d.data(),
+    main_data_size + additional_data_size + num_queries,
+    dimensions,
+    static_cast<int32_t>(10),
+    stream,
+    true,
+    nullptr,
+    nullptr,
+    1.0f,
+    true,
+    -10.0f,
+    10.0f,
+    42ULL,
+    raft::random::GenPC);
 
   // create  dataset DLTensor
   rmm::device_uvector<float> main_d(main_data_size * dimensions, stream);
@@ -321,7 +332,6 @@ TEST(CagraC, BuildExtendSearch)
   cuvsCagraSearch(
     res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter);
 
-  // check neighbors
   ASSERT_TRUE(
     cuvs::devArrMatch(min_cols.data_handle(), neighbors_d.data(), 4, cuvs::Compare<uint32_t>()));
 

ci/validate_wheel.sh

@@ -20,12 +20,13 @@ PYDISTCHECK_ARGS=(
 # PyPI hard limit is 1GiB, but try to keep these as small as possible
 if [[ "${package_dir}" == "python/libcuvs" ]]; then
     if [[ "${RAPIDS_CUDA_MAJOR}" == "12" ]]; then
+        # Cap is below PyPI’s 1 GiB limit; raise when the shipped libcuvs.so grows.
         PYDISTCHECK_ARGS+=(
-            --max-allowed-size-compressed '350Mi'
+            --max-allowed-size-compressed '450Mi'
         )
     else
         PYDISTCHECK_ARGS+=(
-            --max-allowed-size-compressed '220Mi'
+            --max-allowed-size-compressed '270Mi'
         )
     fi
 elif [[ "${package_dir}" != "python/cuvs" ]]; then

cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h

@@ -330,13 +330,6 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index
       std::max(params.graph_degree, params.intermediate_graph_degree);
   }
 
-  nlohmann::json comp_search_conf = collect_conf_with_prefix(conf, "compression_");
-  if (!comp_search_conf.empty()) {
-    auto vpq_pams = params.compression.value_or(cuvs::neighbors::vpq_params{});
-    parse_build_param(comp_search_conf, vpq_pams);
-    params.compression.emplace(vpq_pams);
-  }
-
   if (conf.contains("guarantee_connectivity")) {
     params.guarantee_connectivity = conf.at("guarantee_connectivity");
   }

cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -10,6 +10,8 @@
 
 #include <chrono>
 #include <memory>
+#include <optional>
+#include <variant>
 
 #include "../common/ann_types.hpp"
 #include "../diskann/diskann_wrapper.h"
@@ -165,18 +167,28 @@ void cuvs_cagra_diskann<T, IdxT>::save(const std::string& file) const
 
   // try allocating a buffer for the dataset on host
   try {
-    const cuvs::neighbors::strided_dataset<T, int64_t>* strided_dataset =
-      dynamic_cast<cuvs::neighbors::strided_dataset<T, int64_t>*>(
-        const_cast<cuvs::neighbors::dataset<int64_t>*>(&cagra_build_.get_index()->data()));
-    if (strided_dataset == nullptr) {
-      RAFT_LOG_DEBUG("dynamic_cast to strided_dataset failed");
+    auto const* idx_ptr                                    = cagra_build_.get_index();
+    std::optional<raft::host_matrix<T, int64_t>> h_dataset = std::nullopt;
+    auto const& data_view                                  = idx_ptr->data();
+    if constexpr (cuvs::neighbors::is_padded_dataset_view_v<std::decay_t<decltype(data_view)>>) {
+      auto const& v = data_view;
+      auto n_rows   = v.n_rows();
+      auto dim      = v.dim();
+      auto stride   = v.stride();
+      h_dataset.emplace(raft::make_host_matrix<T, int64_t>(n_rows, dim));
+      raft::copy_matrix(h_dataset->data_handle(),
+                        dim,
+                        v.view().data_handle(),
+                        stride,
+                        dim,
+                        n_rows,
+                        raft::resource::get_cuda_stream(handle_));
     } else {
-      auto h_dataset =
-        raft::make_host_matrix<T, int64_t>(strided_dataset->n_rows(), strided_dataset->dim());
-      raft::copy(h_dataset.data_handle(),
-                 strided_dataset->view().data_handle(),
-                 strided_dataset->n_rows() * strided_dataset->dim(),
-                 raft::resource::get_cuda_stream(handle_));
+      RAFT_LOG_DEBUG("dataset serialization: index dataset is not device_padded_dataset_view");
+    }
+
+    if (h_dataset.has_value()) {
+      raft::resource::sync_stream(handle_);
       std::string dataset_base_file = file + ".data";
       std::ofstream dataset_of(dataset_base_file, std::ios::out | std::ios::binary);
       if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_base_file.c_str()); }
@@ -187,7 +199,7 @@ void cuvs_cagra_diskann<T, IdxT>::save(const std::string& file) const
       dataset_of.write((char*)&size, sizeof(int));
       dataset_of.write((char*)&dim, sizeof(int));
       for (int i = 0; i < size; i++) {
-        dataset_of.write((char*)(h_dataset.data_handle() + i * h_dataset.extent(1)),
+        dataset_of.write((char*)(h_dataset->data_handle() + i * h_dataset->extent(1)),
                          dim * sizeof(T));
       }
       dataset_of.close();

cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h

@@ -82,11 +82,7 @@ class cuvs_cagra_hnswlib : public algo<T>, public algo_gpu {
 template <typename T, typename IdxT>
 void cuvs_cagra_hnswlib<T, IdxT>::build(const T* dataset, size_t nrow)
 {
-  // when the data set is on host, we can pass it directly to HNSW
-  bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1;
-
   auto dataset_view = raft::make_host_matrix_view<const T, int64_t>(dataset, nrow, this->dim_);
-  // convert the index to HNSW format
   hnsw_index_ = cuvs::neighbors::hnsw::build(handle_, build_param_.hnsw_index_params, dataset_view);
 }