Recover cagra-q path in C API and downstream language wrappers

c/include/cuvs/core/dataset.h

@@ -20,7 +20,9 @@ extern "C" {
  */
 typedef enum {
   CUVS_DATASET_LAYOUT_STANDARD = 0,
-  CUVS_DATASET_LAYOUT_PADDED   = 1
+  CUVS_DATASET_LAYOUT_PADDED   = 1,
+  /** Device VPQ storage with f16 codebooks (CAGRA-Q search dataset). */
+  CUVS_DATASET_LAYOUT_VPQ_F16  = 2
 } cuvsDatasetLayout_t;
 
 /**

c/include/cuvs/neighbors/cagra.h

@@ -255,6 +255,24 @@ CUVS_EXPORT cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionPar
  */
 CUVS_EXPORT cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params);
 
+/**
+ * @brief Train an owning device VPQ (f16 codebook) dataset from a device-padded source.
+ *
+ * Used for CAGRA-Q: build a dense CAGRA index, train VPQ with this factory, then attach via
+ * `cuvsCagraUpdateDataset`. Caller owns the returned dataset and must keep it alive while any
+ * index uses it. Metric for subsequent search must remain `L2Expanded`.
+ *
+ * @param[in] res cuvs resources
+ * @param[in] source_dataset device-padded dataset (owning or view)
+ * @param[in] params VPQ compression params; NULL selects defaults
+ * @param[out] vpq_dataset newly allocated owning VPQ dataset handle
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsDatasetMakeVpq(cuvsResources_t res,
+                                           cuvsDataset_t source_dataset,
+                                           cuvsCagraCompressionParams_t params,
+                                           cuvsDataset_t* vpq_dataset);
+
 /**
  * @brief Allocate ACE params, and populate with default values
  *
@@ -580,21 +598,25 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, DLManag
 CUVS_EXPORT cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, DLManagedTensor* graph);
 
 /**
- * @brief Update a CAGRA index with a device-padded dataset.
+ * @brief Update a CAGRA index with a device dataset (padded or VPQ).
+ *
+ * This is the centralized dataset update/attach operation for C callers.
+ *
+ * - Device-padded dataset: if \p index is already device-padded, its dataset view is replaced in
+ *   place (same index object); otherwise the index is converted via attach and rebound.
+ * - Device VPQ_F16 dataset (from `cuvsDatasetMakeVpq`): if \p index is already VPQ-typed, its
+ *   dataset view is replaced in place; otherwise the graph is copied into a new VPQ-typed index
+ *   (CAGRA-Q). Search requires metric `L2Expanded`. The VPQ handle must be owning.
  *
- * This is the centralized dataset update operation for C callers. If \p index
- * is already device-padded, its dataset view is replaced in place. Otherwise,
- * the index is converted and its opaque handle is rebound to a search-ready
- * device-padded index. Caller retains ownership of
- * \p device_padded_dataset and must keep it alive while \p index uses it.
+ * Caller retains ownership of \p dataset and must keep it alive while \p index uses it.
  *
- * @param[in] res             cuvsResources_t opaque C handle
- * @param[in] device_padded_dataset owning or non-owning device-padded dataset handle
- * @param[inout] index        CAGRA index handle
+ * @param[in] res      cuvsResources_t opaque C handle
+ * @param[in] dataset  device-padded or owning device VPQ_F16 dataset handle
+ * @param[inout] index CAGRA index handle
  * @return cuvsError_t
  */
 CUVS_EXPORT cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res,
-                                               cuvsDataset_t device_padded_dataset,
+                                               cuvsDataset_t dataset,
                                                cuvsCagraIndex_t index);
 
 /**

c/tests/neighbors/ann_cagra_c.cu

@@ -19,6 +19,7 @@
 #include <string>
 #include <type_traits>
 #include <unistd.h>
+#include <vector>
 
 #include <cuda_runtime.h>
 #include <gtest/gtest.h>
@@ -2008,3 +2009,109 @@ TEST(CagraC, SearchMultiPartitionMultiKernelRejected)
   }
   cuvsResourcesDestroy(res);
 }
+
+TEST(CagraC, BuildAttachVpqSearch)
+{
+  // CAGRA-Q smoke test: dense build → MakeVpq → UpdateDataset(VPQ) → Search.
+  constexpr int64_t n_rows    = 256;
+  constexpr int64_t dim       = 32;
+  constexpr int64_t n_queries = 4;
+  constexpr int64_t k         = 1;
+
+  cuvsResources_t res;
+  ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS);
+  cudaStream_t stream;
+  ASSERT_EQ(cuvsStreamGet(res, &stream), CUVS_SUCCESS);
+
+  rmm::device_uvector<float> dataset_d(n_rows * dim, stream);
+  {
+    std::vector<float> host(n_rows * dim);
+    for (int64_t i = 0; i < n_rows * dim; ++i) {
+      host[i] = static_cast<float>((i % 17) + 1);
+    }
+    raft::copy(dataset_d.data(), host.data(), host.size(), stream);
+  }
+
+  // dim=32 float already matches CAGRA padded row width; MakePadded refuses a
+  // no-op device copy — wrap with MakePaddedView instead.
+  DLManagedTensor dataset_tensor{};
+  dataset_tensor.dl_tensor.data               = dataset_d.data();
+  dataset_tensor.dl_tensor.device.device_type = kDLCUDA;
+  dataset_tensor.dl_tensor.ndim               = 2;
+  dataset_tensor.dl_tensor.dtype              = {kDLFloat, 32, 1};
+  int64_t dataset_shape[2]                    = {n_rows, dim};
+  dataset_tensor.dl_tensor.shape              = dataset_shape;
+  dataset_tensor.dl_tensor.strides            = nullptr;
+
+  cuvsDataset_t padded;
+  ASSERT_EQ(cuvsDatasetMakePaddedView(res, &dataset_tensor, &padded), CUVS_SUCCESS);
+
+  cuvsCagraIndexParams_t build_params;
+  ASSERT_EQ(cuvsCagraIndexParamsCreate(&build_params), CUVS_SUCCESS);
+  cuvsCagraIndex_t index;
+  ASSERT_EQ(cuvsCagraIndexCreate(&index), CUVS_SUCCESS);
+  ASSERT_EQ(cuvsCagraBuild(res, build_params, padded, index), CUVS_SUCCESS);
+
+  cuvsCagraCompressionParams_t compression;
+  ASSERT_EQ(cuvsCagraCompressionParamsCreate(&compression), CUVS_SUCCESS);
+  compression->pq_bits = 8;
+  compression->pq_dim  = 8;
+
+  cuvsDataset_t vpq = nullptr;
+  ASSERT_EQ(cuvsDatasetMakeVpq(res, padded, compression, &vpq), CUVS_SUCCESS);
+  {
+    cuvsDatasetLayout_t layout;
+    ASSERT_EQ(cuvsDatasetGetLayout(vpq, &layout), CUVS_SUCCESS);
+    EXPECT_EQ(layout, CUVS_DATASET_LAYOUT_VPQ_F16);
+    bool owning = false;
+    ASSERT_EQ(cuvsDatasetGetIsOwning(vpq, &owning), CUVS_SUCCESS);
+    EXPECT_TRUE(owning);
+  }
+
+  ASSERT_EQ(cuvsCagraUpdateDataset(res, vpq, index), CUVS_SUCCESS);
+
+  rmm::device_uvector<float> queries_d(n_queries * dim, stream);
+  raft::copy(queries_d.data(), dataset_d.data(), n_queries * dim, stream);
+  DLManagedTensor queries_tensor{};
+  queries_tensor.dl_tensor.data               = queries_d.data();
+  queries_tensor.dl_tensor.device.device_type = kDLCUDA;
+  queries_tensor.dl_tensor.ndim               = 2;
+  queries_tensor.dl_tensor.dtype              = {kDLFloat, 32, 1};
+  int64_t queries_shape[2]                    = {n_queries, dim};
+  queries_tensor.dl_tensor.shape              = queries_shape;
+
+  rmm::device_uvector<uint32_t> neighbors_d(n_queries * k, stream);
+  DLManagedTensor neighbors_tensor{};
+  neighbors_tensor.dl_tensor.data               = neighbors_d.data();
+  neighbors_tensor.dl_tensor.device.device_type = kDLCUDA;
+  neighbors_tensor.dl_tensor.ndim               = 2;
+  neighbors_tensor.dl_tensor.dtype              = {kDLUInt, 32, 1};
+  int64_t neighbors_shape[2]                    = {n_queries, k};
+  neighbors_tensor.dl_tensor.shape              = neighbors_shape;
+
+  rmm::device_uvector<float> distances_d(n_queries * k, stream);
+  DLManagedTensor distances_tensor{};
+  distances_tensor.dl_tensor.data               = distances_d.data();
+  distances_tensor.dl_tensor.device.device_type = kDLCUDA;
+  distances_tensor.dl_tensor.ndim               = 2;
+  distances_tensor.dl_tensor.dtype              = {kDLFloat, 32, 1};
+  int64_t distances_shape[2]                    = {n_queries, k};
+  distances_tensor.dl_tensor.shape              = distances_shape;
+
+  cuvsFilter filter;
+  filter.type = NO_FILTER;
+  filter.addr = (uintptr_t)NULL;
+  cuvsCagraSearchParams_t search_params;
+  ASSERT_EQ(cuvsCagraSearchParamsCreate(&search_params), CUVS_SUCCESS);
+  ASSERT_EQ(cuvsCagraSearch(
+              res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter),
+            CUVS_SUCCESS);
+
+  cuvsCagraSearchParamsDestroy(search_params);
+  cuvsCagraCompressionParamsDestroy(compression);
+  cuvsDatasetDestroy(vpq);
+  cuvsCagraIndexDestroy(index);
+  cuvsCagraIndexParamsDestroy(build_params);
+  cuvsDatasetDestroy(padded);
+  cuvsResourcesDestroy(res);
+}

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

@@ -400,11 +400,12 @@ void cuvs_cagra<T, IdxT>::compress_dataset(const T* dataset, size_t nrow)
                "cagra: compression_* (CAGRA-Q) requires the graph in memory; it cannot be combined "
                "with a disk-resident (ACE) graph.");
   auto rows = static_cast<int64_t>(nrow);
-  // make_vpq_dataset() reads the rows wherever they are: host-resident ones are subsampled and
-  // encoded in bounded batches instead of being staged on the device.
+  // make_device_vpq_dataset() reads the rows wherever they are: host-resident ones are subsampled
+  // and encoded in bounded batches instead of being staged on the device.
   auto src = raft::make_device_matrix_view<const T, int64_t, raft::row_major>(dataset, rows, dim_);
   vpq_dataset_ = std::make_shared<cuvs::neighbors::device_vpq_dataset<half, int64_t>>(
-    cuvs::preprocessing::quantize::pq::make_vpq_dataset(handle_, *index_params_.compression, src));
+    cuvs::preprocessing::quantize::pq::make_device_vpq_dataset(
+      handle_, *index_params_.compression, src));
   vpq_index_ = std::make_shared<cuvs::neighbors::cagra::vpq_f16_index<T, IdxT>>(
     handle_, parse_metric_type(metric_), vpq_dataset_->as_dataset_view(), index_->graph());
 

cpp/include/cuvs/neighbors/cagra.hpp

@@ -4548,6 +4548,42 @@ auto convert_standard_to_padded_index(
   return out;
 }
 
+/**
+ * @brief Convert a dense (non-VPQ) CAGRA index into a VPQ f16 index and attach a VPQ dataset.
+ *
+ * Copies graph/source-indices from `src` into a new `vpq_f16_index` and attaches `vpq_dataset`.
+ * Caller owns `vpq_dataset` storage for the lifetime of the returned index.
+ */
+template <typename T, typename IdxT, typename IndexViewT>
+  requires cuvs::neighbors::ann_dataset_view<IndexViewT>
+auto convert_dense_to_vpq_f16_index(
+  raft::resources const& res,
+  index<T, IdxT, IndexViewT> const& src,
+  cuvs::neighbors::device_vpq_dataset_view<half, int64_t> const& vpq_dataset)
+  -> vpq_f16_index<T, IdxT>
+{
+  RAFT_EXPECTS(vpq_dataset.n_rows() == src.size(),
+               "VPQ dataset row count must match the index size");
+
+  vpq_f16_index<T, IdxT> out(res, src.metric());
+  if (src.graph().extent(0) > 0) {
+    using GraphIndexType = typename index<T, IdxT, IndexViewT>::graph_index_type;
+    auto graph_host =
+      raft::make_host_matrix<GraphIndexType, int64_t>(src.graph().extent(0), src.graph().extent(1));
+    raft::copy(graph_host.data_handle(),
+               src.graph().data_handle(),
+               src.graph().size(),
+               raft::resource::get_cuda_stream(res));
+    raft::resource::sync_stream(res);
+    out.update_graph(res, raft::make_const_mdspan(graph_host.view()));
+  }
+  if (src.source_indices().has_value()) {
+    out.update_source_indices(res, src.source_indices().value());
+  }
+  out.update_device_dataset_same_layout(res, vpq_dataset);
+  return out;
+}
+
 /**
  * @brief Attach a device-padded dataset and return a search-ready padded-device index.
  *
@@ -4598,6 +4634,35 @@ auto attach_dataset(
   }
 }
 
+/**
+ * @brief Attach a device VPQ (f16 codebook) dataset and return a search-ready `vpq_f16_index`.
+ *
+ * Builds a new VPQ-typed index by copying the graph from `idx`. Caller owns `vpq_dataset` storage.
+ *
+ * If `idx` is already a `vpq_f16_index`, call `idx.update_device_dataset_same_layout(res,
+ * vpq_dataset)` directly.
+ */
+template <typename T, typename IdxT, typename IndexViewT>
+  requires cuvs::neighbors::ann_dataset_view<IndexViewT>
+auto attach_dataset(raft::resources const& res,
+                    index<T, IdxT, IndexViewT> const& idx,
+                    cuvs::neighbors::device_vpq_dataset_view<half, int64_t> const& vpq_dataset)
+  -> vpq_f16_index<T, IdxT>
+{
+  if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v<IndexViewT>) {
+    RAFT_LOG_WARN(
+      "cagra::attach_dataset called with an already vpq_f16 index. "
+      "To avoid an unnecessary index copy, call "
+      "index.update_device_dataset_same_layout(res, vpq_dataset) "
+      "directly on the original index.");
+    RAFT_FAIL(
+      "cagra::attach_dataset: vpq_f16_index input is not supported in this overload. "
+      "Call index.update_device_dataset_same_layout(res, vpq_dataset) directly.");
+  } else {
+    return convert_dense_to_vpq_f16_index(res, idx, vpq_dataset);
+  }
+}
+
 }  // namespace cagra
 }  // namespace neighbors
 }  // namespace CUVS_EXPORT cuvs

cpp/include/cuvs/preprocessing/quantize/pq.hpp

@@ -281,24 +281,25 @@ namespace detail {
  *
  * Typical **CAGRA** usage: build the graph on dense vectors, then attach VPQ for search (metric
  * must remain `L2Expanded` for this path). Train VPQ from the same CAGRA-padded device layout you
- * used for graph build, keep the `device_vpq_dataset` alive, and call
- * `index::update_device_dataset_same_layout` with a non-owning view.
+ * used for graph build, keep the `device_vpq_dataset` alive, and attach it with
+ * `cagra::attach_dataset` (returns a `vpq_f16_index`).
  *
  * @code{.cpp}
  * #include <cuvs/neighbors/cagra.hpp>
  * #include <cuvs/preprocessing/quantize/pq.hpp>
  *
- * // `idx` is a `cagra::index<float, uint32_t>` with graph built on dense rows.
+ * // `idx` is a dense CAGRA index with graph built on padded rows.
  * // `padded` is a `device_padded_dataset_view<float, int64_t>` view of those same rows.
  * cuvs::neighbors::vpq_params vpq_params{};
- * auto vpq = cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, padded);
+ * auto vpq = cuvs::preprocessing::quantize::pq::make_device_vpq_dataset(
+ *   res, vpq_params, padded);
  * idx.update_device_dataset_same_layout(res, vpq.as_dataset_view());
  * @endcode
  */
 template <typename SrcT>
-[[nodiscard]] auto make_vpq_dataset(raft::resources const& res,
-                                    cuvs::neighbors::vpq_params const& params,
-                                    SrcT const& src)
+[[nodiscard]] auto make_device_vpq_dataset(raft::resources const& res,
+                                           cuvs::neighbors::vpq_params const& params,
+                                           SrcT const& src)
   -> cuvs::neighbors::device_vpq_dataset<half, int64_t>
 {
   // A cuVS dataset keeps its logical width in `dim()` while `view()` spans the full row pitch.
@@ -310,7 +311,7 @@ template <typename SrcT>
     auto const rows    = src.view();
     using value_type   = typename decltype(rows)::value_type;
     using extents_type = raft::matrix_extent<int64_t>;
-    return make_vpq_dataset(
+    return make_device_vpq_dataset(
       res,
       params,
       raft::mdspan<const value_type, extents_type, raft::layout_stride>{
@@ -321,11 +322,11 @@ template <typename SrcT>
     using value_type = typename SrcT::value_type;
     static_assert(std::is_same_v<value_type, float> || std::is_same_v<value_type, half> ||
                     std::is_same_v<value_type, int8_t> || std::is_same_v<value_type, uint8_t>,
-                  "make_vpq_dataset: element type must be float, half, int8_t or uint8_t");
+                  "make_device_vpq_dataset: element type must be float, half, int8_t or uint8_t");
     const int64_t n_rows = src.extent(0);
     const int64_t dim    = src.extent(1);
     const int64_t stride = src.stride(0) > 0 ? src.stride(0) : dim;
-    RAFT_EXPECTS(n_rows > 0, "make_vpq_dataset: dataset is empty");
+    RAFT_EXPECTS(n_rows > 0, "make_device_vpq_dataset: dataset is empty");
     return detail::vpq_train_from_rows(
       res, params, src.data_handle(), raft::get_cuda_data_type<value_type>(), n_rows, dim, stride);
   }

cpp/src/preprocessing/quantize/pq.cu

@@ -5,6 +5,8 @@
 
 #include "./detail/pq.cuh"
 
+#include <cuvs/core/cuda_fp16.hpp>
+#include <cuvs/core/export.hpp>
 #include <cuvs/preprocessing/quantize/pq.hpp>
 
 #include <raft/matrix/copy.cuh>
@@ -92,7 +94,7 @@ auto train_from_rows(raft::resources const& res,
   if (device_ptr == nullptr) {
     // A host mdspan makes training subsample the rows and encoding stream them in bounded batches,
     // so the dense dataset is never staged on the device.
-    RAFT_EXPECTS(stride == dim, "make_vpq_dataset: host input must be tightly packed");
+    RAFT_EXPECTS(stride == dim, "make_device_vpq_dataset: host input must be tightly packed");
     auto row_view = raft::make_host_matrix_view<const T, int64_t>(src_ptr, n_rows, dim);
     return detail::vpq_build_half(res, params, row_view);
   }
@@ -132,7 +134,8 @@ auto vpq_train_from_rows(raft::resources const& res,
       return train_from_rows(
         res, params, static_cast<uint8_t const*>(src_ptr), n_rows, dim, stride);
     default:
-      RAFT_FAIL("make_vpq_dataset: unsupported dataset element type %d", static_cast<int>(dtype));
+      RAFT_FAIL("make_device_vpq_dataset: unsupported dataset element type %d",
+                static_cast<int>(dtype));
   }
 }