Add cagra.extend bindings for python

cpp/include/cuvs/neighbors/cagra.h

@@ -525,14 +525,12 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res,
  * @param[in] params cuvsCagraExtendParams_t used to extend CAGRA index
  * @param[in] additional_dataset DLManagedTensor* additional dataset
  * @param[in,out] index cuvsCagraIndex_t CAGRA index
- * @param[out] return_dataset DLManagedTensor* extended dataset
  * @return cuvsError_t
  */
 cuvsError_t cuvsCagraExtend(cuvsResources_t res,
                             cuvsCagraExtendParams_t params,
                             DLManagedTensor* additional_dataset,
-                            cuvsCagraIndex_t index,
-                            DLManagedTensor* return_dataset);
+                            cuvsCagraIndex_t index);
 
 /**
  * @}

cpp/src/neighbors/cagra_c.cpp

@@ -164,32 +164,24 @@ template <typename T>
 void _extend(cuvsResources_t res,
              cuvsCagraExtendParams params,
              cuvsCagraIndex index,
-             DLManagedTensor* additional_dataset_tensor,
-             DLManagedTensor* return_tensor)
+             DLManagedTensor* additional_dataset_tensor)
 {
-  auto dataset          = additional_dataset_tensor->dl_tensor;
-  auto return_dl_tensor = return_tensor->dl_tensor;
-  auto index_ptr        = reinterpret_cast<cuvs::neighbors::cagra::index<T, uint32_t>*>(index.addr);
-  auto res_ptr          = reinterpret_cast<raft::resources*>(res);
+  auto dataset   = additional_dataset_tensor->dl_tensor;
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<T, uint32_t>*>(index.addr);
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
 
   // TODO: use C struct here (see issue #487)
   auto extend_params           = cuvs::neighbors::cagra::extend_params();
   extend_params.max_chunk_size = params.max_chunk_size;
 
-  if (cuvs::core::is_dlpack_device_compatible(dataset) &&
-      cuvs::core::is_dlpack_device_compatible(return_dl_tensor)) {
-    using mdspan_type        = raft::device_matrix_view<T const, int64_t, raft::row_major>;
-    using mdspan_return_type = raft::device_matrix_view<T, int64_t, raft::row_major>;
-    auto mds                 = cuvs::core::from_dlpack<mdspan_type>(additional_dataset_tensor);
-    auto return_mds          = cuvs::core::from_dlpack<mdspan_return_type>(return_tensor);
-    cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr, return_mds);
-  } else if (cuvs::core::is_dlpack_host_compatible(dataset) &&
-             cuvs::core::is_dlpack_host_compatible(return_dl_tensor)) {
-    using mdspan_type        = raft::host_matrix_view<T const, int64_t, raft::row_major>;
-    using mdspan_return_type = raft::device_matrix_view<T, int64_t, raft::row_major>;
-    auto mds                 = cuvs::core::from_dlpack<mdspan_type>(additional_dataset_tensor);
-    auto return_mds          = cuvs::core::from_dlpack<mdspan_return_type>(return_tensor);
-    cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr, return_mds);
+  if (cuvs::core::is_dlpack_device_compatible(dataset)) {
+    using mdspan_type = raft::device_matrix_view<T const, int64_t, raft::row_major>;
+    auto mds          = cuvs::core::from_dlpack<mdspan_type>(additional_dataset_tensor);
+    cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr);
+  } else if (cuvs::core::is_dlpack_host_compatible(dataset)) {
+    using mdspan_type = raft::host_matrix_view<T const, int64_t, raft::row_major>;
+    auto mds          = cuvs::core::from_dlpack<mdspan_type>(additional_dataset_tensor);
+    cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr);
   } else {
     RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
               dataset.dtype.code,
@@ -540,19 +532,18 @@ extern "C" cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res,
 extern "C" cuvsError_t cuvsCagraExtend(cuvsResources_t res,
                                        cuvsCagraExtendParams_t params,
                                        DLManagedTensor* additional_dataset_tensor,
-                                       cuvsCagraIndex_t index_c_ptr,
-                                       DLManagedTensor* return_dataset_tensor)
+                                       cuvsCagraIndex_t index_c_ptr)
 {
   return cuvs::core::translate_exceptions([=] {
     auto dataset = additional_dataset_tensor->dl_tensor;
     auto index   = *index_c_ptr;
 
     if ((dataset.dtype.code == kDLFloat) && (dataset.dtype.bits == 32)) {
-      _extend<float>(res, *params, index, additional_dataset_tensor, return_dataset_tensor);
+      _extend<float>(res, *params, index, additional_dataset_tensor);
     } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) {
-      _extend<int8_t>(res, *params, index, additional_dataset_tensor, return_dataset_tensor);
+      _extend<int8_t>(res, *params, index, additional_dataset_tensor);
     } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) {
-      _extend<uint8_t>(res, *params, index, additional_dataset_tensor, return_dataset_tensor);
+      _extend<uint8_t>(res, *params, index, additional_dataset_tensor);
     } else {
       RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
                 dataset.dtype.code,

cpp/tests/neighbors/ann_cagra_c.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -210,20 +210,6 @@ TEST(CagraC, BuildExtendSearch)
   additional_dataset_tensor.dl_tensor.shape              = additional_dataset_shape;
   additional_dataset_tensor.dl_tensor.strides            = nullptr;
 
-  // create tensor for that points to the extended dataset
-  rmm::device_uvector<float> extend_return_d((additional_data_size + main_data_size) * dimensions,
-                                             stream);
-  DLManagedTensor additional_dataset_return_tensor;
-  additional_dataset_return_tensor.dl_tensor.data               = extend_return_d.data();
-  additional_dataset_return_tensor.dl_tensor.device.device_type = kDLCUDA;
-  additional_dataset_return_tensor.dl_tensor.ndim               = 2;
-  additional_dataset_return_tensor.dl_tensor.dtype.code         = kDLFloat;
-  additional_dataset_return_tensor.dl_tensor.dtype.bits         = 32;
-  additional_dataset_return_tensor.dl_tensor.dtype.lanes        = 1;
-  int64_t additional_return_dataset_shape[2] = {additional_data_size + main_data_size, dimensions};
-  additional_dataset_return_tensor.dl_tensor.shape   = additional_return_dataset_shape;
-  additional_dataset_return_tensor.dl_tensor.strides = nullptr;
-
   // create index
   cuvsCagraIndex_t index;
   cuvsCagraIndexCreate(&index);
@@ -238,8 +224,7 @@ TEST(CagraC, BuildExtendSearch)
   // extend index
   cuvsCagraExtendParams_t extend_params;
   cuvsCagraExtendParamsCreate(&extend_params);
-  cuvsCagraExtend(
-    res, extend_params, &additional_dataset_tensor, index, &additional_dataset_return_tensor);
+  cuvsCagraExtend(res, extend_params, &additional_dataset_tensor, index);
 
   // create queries DLTensor
   rmm::device_uvector<float> queries_d(num_queries * dimensions, stream);
@@ -342,15 +327,6 @@ TEST(CagraC, BuildExtendSearch)
   cuvsCagraSearch(
     res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter);
 
-  // make sure that extend_return_d points to the extended dataset
-  ASSERT_TRUE(cuvs::devArrMatch(
-    main_d.data(), extend_return_d.data(), main_d.size(), cuvs::Compare<float>()));
-
-  ASSERT_TRUE(cuvs::devArrMatch(additional_d.data(),
-                                extend_return_d.data() + main_d.size(),
-                                additional_d.size(),
-                                cuvs::Compare<float>()));
-
   // check neighbors
   ASSERT_TRUE(
     cuvs::devArrMatch(min_cols.data_handle(), neighbors_d.data(), 4, cuvs::Compare<uint32_t>()));

go/cagra/cagra.go

@@ -51,13 +51,12 @@ func BuildIndex[T any](Resources cuvs.Resource, params *IndexParams, dataset *cu
 // * `Resources` - Resources to use
 // * `params` - Parameters for extending the index
 // * `additional_dataset` - A row-major Tensor on the device to extend the index with
-// * `return_dataset` - A row-major Tensor on the device that will receive the extended dataset
 // * `index` - CagraIndex to extend
-func ExtendIndex[T any](Resources cuvs.Resource, params *ExtendParams, additional_dataset *cuvs.Tensor[T], return_dataset *cuvs.Tensor[T], index *CagraIndex) error {
+func ExtendIndex[T any](Resources cuvs.Resource, params *ExtendParams, additional_dataset *cuvs.Tensor[T], index *CagraIndex) error {
 	if !index.trained {
 		return errors.New("index needs to be built before calling extend")
 	}
-	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraExtend(C.ulong(Resources.Resource), params.params, (*C.DLManagedTensor)(unsafe.Pointer(additional_dataset.C_tensor)), index.index, (*C.DLManagedTensor)(unsafe.Pointer(return_dataset.C_tensor)))))
+	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraExtend(C.ulong(Resources.Resource), params.params, (*C.DLManagedTensor)(unsafe.Pointer(additional_dataset.C_tensor)), index.index)))
 	if err != nil {
 		return err
 	}

python/cuvs/cuvs/neighbors/cagra/__init__.py

@@ -15,10 +15,12 @@
 
 from .cagra import (
     CompressionParams,
+    ExtendParams,
     Index,
     IndexParams,
     SearchParams,
     build,
+    extend,
     from_graph,
     load,
     save,
@@ -27,10 +29,12 @@
 
 __all__ = [
     "CompressionParams",
+    "ExtendParams",
     "Index",
     "IndexParams",
     "SearchParams",
     "build",
+    "extend",
     "from_graph",
     "load",
     "save",

python/cuvs/cuvs/neighbors/cagra/cagra.pxd

@@ -175,6 +175,19 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil:
                                        DLManagedTensor * dataset,
                                        cuvsCagraIndex_t index)
 
+    ctypedef struct cuvsCagraExtendParams:
+        uint32_t max_chunk_size
+
+    ctypedef cuvsCagraExtendParams* cuvsCagraExtendParams_t
+
+    cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params)
+    cuvsError_t cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params)
+    cuvsError_t cuvsCagraExtend(cuvsResources_t res,
+                                cuvsCagraExtendParams_t params,
+                                DLManagedTensor* additional_dataset,
+                                cuvsCagraIndex_t index)
+
+
 cdef class Index:
     """
     CAGRA index object. This object stores the trained CAGRA index state

python/cuvs/cuvs/neighbors/cagra/cagra.pyx

@@ -350,7 +350,7 @@ def build(IndexParams index_params, dataset, resources=None):
     ----------
     index_params : IndexParams object
     dataset : CUDA array interface compliant matrix shape (n_samples, dim)
-        Supported dtype [float, int8, uint8]
+        Supported dtype [float, half, int8, uint8]
     {resources_docstring}
 
     Returns
@@ -841,3 +841,73 @@ def from_graph(graph, dataset, metric="sqeuclidean", resources=None):
         idx.index))
     idx.trained = True
     return idx
+
+
+cdef class ExtendParams:
+    """ Supplemental parameters to extend CAGRA Index
+
+    Parameters
+    ----------
+    max_chunk_size : int
+        The additional dataset is divided into chunks and added to the graph.
+        This is the knob to adjust the tradeoff between the recall and
+        operation throughput. Large chunk sizes can result in high throughput,
+        but use more working memory (O(max_chunk_size*degree^2)). This can also
+        degrade recall because no edges are added between the nodes in the same
+        chunk. Auto select when 0.
+    """
+
+    cdef cuvsCagraExtendParams* params
+
+    def __cinit__(self):
+        check_cuvs(cuvsCagraExtendParamsCreate(&self.params))
+
+    def __dealloc__(self):
+        check_cuvs(cuvsCagraExtendParamsDestroy(self.params))
+
+    def __init__(self, *,
+                 max_chunk_size=None):
+        if max_chunk_size is not None:
+            self.params.max_chunk_size = max_chunk_size
+
+    @property
+    def max_chunk_size(self):
+        return self.params.max_chunk_size
+
+
+@auto_sync_resources
+def extend(ExtendParams params, Index index, additional_dataset,
+           resources=None):
+    """
+    Extend a CAGRA index with additional vectors
+
+    Parameters
+    ----------
+    params : ExtendParams object
+    index: Index
+       Existing cagra index to extend
+    additional_dataset : CUDA array interface compliant matrix shape
+        Supported dtype [float, half, int8, uint8]
+    {resources_docstring}
+
+    """
+    dataset_ai = wrap_array(additional_dataset)
+    _check_input_array(dataset_ai, [np.dtype("float32"),
+                                    np.dtype("float16"),
+                                    np.dtype("byte"),
+                                    np.dtype("ubyte")])
+
+    cdef cydlpack.DLManagedTensor* dataset_dlpack = \
+        cydlpack.dlpack_c(dataset_ai)
+
+    cdef cuvsResources_t res = <cuvsResources_t>resources.get_c_obj()
+
+    with cuda_interruptible():
+        check_cuvs(cuvsCagraExtend(
+            res,
+            params.params,
+            dataset_dlpack,
+            index.index
+        ))
+
+    return index

python/cuvs/cuvs/tests/test_cagra.py

@@ -41,7 +41,7 @@ def run_cagra_build_search_test(
     array_type="device",
     compare=True,
     inplace=True,
-    add_data_on_build=True,
+    test_extend=False,
     search_params={},
     compression=None,
 ):
@@ -60,26 +60,24 @@ def run_cagra_build_search_test(
         compression=compression,
     )
 
-    if array_type == "device":
-        index = cagra.build(build_params, dataset_device)
-    else:
-        index = cagra.build(build_params, dataset)
-
-    if not add_data_on_build:
+    if test_extend:
         dataset_1 = dataset[: n_rows // 2, :]
         dataset_2 = dataset[n_rows // 2 :, :]
-        indices_1 = np.arange(n_rows // 2, dtype=np.uint32)
-        indices_2 = np.arange(n_rows // 2, n_rows, dtype=np.uint32)
+        extend_params = cagra.ExtendParams()
         if array_type == "device":
             dataset_1_device = device_ndarray(dataset_1)
             dataset_2_device = device_ndarray(dataset_2)
-            indices_1_device = device_ndarray(indices_1)
-            indices_2_device = device_ndarray(indices_2)
-            index = cagra.extend(index, dataset_1_device, indices_1_device)
-            index = cagra.extend(index, dataset_2_device, indices_2_device)
+
+            index = cagra.build(build_params, dataset_1_device)
+            index = cagra.extend(extend_params, index, dataset_2_device)
+        else:
+            index = cagra.build(build_params, dataset_1)
+            index = cagra.extend(index, dataset_2)
+    else:
+        if array_type == "device":
+            index = cagra.build(build_params, dataset_device)
         else:
-            index = cagra.extend(index, dataset_1, indices_1)
-            index = cagra.extend(index, dataset_2, indices_2)
+            index = cagra.build(build_params, dataset)
 
     queries = generate_data((n_queries, n_cols), dtype)
     out_idx = np.zeros((n_queries, k), dtype=np.uint32)
@@ -185,23 +183,23 @@ def test_filtered_cagra(sparsity):
         {
             "intermediate_graph_degree": 64,
             "graph_degree": 32,
-            "add_data_on_build": True,
+            "test_extend": False,
             "k": 1,
             "metric": "sqeuclidean",
             "build_algo": "ivf_pq",
         },
         {
             "intermediate_graph_degree": 32,
             "graph_degree": 16,
-            "add_data_on_build": False,
+            "test_extend": True,
             "k": 5,
             "metric": "sqeuclidean",
             "build_algo": "ivf_pq",
         },
         {
             "intermediate_graph_degree": 128,
             "graph_degree": 32,
-            "add_data_on_build": True,
+            "test_extend": False,
             "k": 10,
             "metric": "inner_product",
             "build_algo": "nn_descent",
@@ -212,6 +210,7 @@ def test_cagra_index_params(params):
     # Note that inner_product tests use normalized input which we cannot
     # represent in int8, therefore we test only sqeuclidean metric here.
     run_cagra_build_search_test(
+        test_extend=params["test_extend"],
         k=params["k"],
         metric=params["metric"],
         graph_degree=params["graph_degree"],