`update_dataset` CAGRA by aamijar · Pull Request #2427 · NVIDIA/cuvs

aamijar · 2026-08-08T01:17:22Z

This PR introduces the update_dataset function in the cpp api.
update_dataset should swap out the dataset from an index. For instance, from a device_standard_dataset_view -> device_padded_dataset_view.

The major use case here is when a user builds a cagra index with a certain type of dataset. But during search they want to use a different type of dataset. In the future this may look like device_standard_dataset_view -> device_pq_dataset_view.

The type of dataset is coupled to the index type. So we need to construct a new index object altogether. To avoid copying expensive member variables like the cagra graph or source indices we simply move them to the new object instead. This means the user relinquishes the old index object.

A new constructor in cagra.hpp has been introduced to move and take ownership of the resources from an existing cagra index and assign the new dataset.

We intend to remove the attach_dataset and other helper functions such as convert_standard_to_padded_index and convert_host_to_device_index. The callers that use the removed functions should be updated to use update_dataset instead.

Resolves #2404

I've annotated this PR below to make it easier to review.

copy-pr-bot · 2026-08-08T01:17:25Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

aamijar · 2026-08-12T03:57:24Z

-#include <numeric>
 #include <optional>
 #include <string>
+#include <type_traits>


Cleaning up some unused headers and reordering.

aamijar · 2026-08-12T03:57:56Z

-          typename IdxT,
-          cuvs::neighbors::ann_dataset_view DatasetViewT =
-            cuvs::neighbors::device_padded_dataset_view<T, int64_t>>
-struct index;


Removing this unnecessary forward declaration.

aamijar · 2026-08-12T03:59:34Z

-  index(raft::resources const& res,
-        cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded)
+  explicit index(raft::resources const& res,
+                 cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded)


Adding explicit keyword to avoid unintentional implicit conversions.

aamijar · 2026-08-12T04:02:33Z

+      graph_degree_(other.graph_degree_)
+  {
+    update_dataset(res, dataset);
+  }


The core implementation of the free function update_dataset is in this constructor. The update_dataset is really a wrapper around this.
This constructor creates a cagra index and takes ownership of the resources of the old one (specifically other.graph_ and other.source_indices_).
We reuse the member function index.update_dataset() to avoid duplicating code.

The member function index.update_dataset() is compatible with the same type only.
The update_dataset free function is for updating datasets of different types (ex. standard to padded).

Arent we also supposed to move the file descriptors for ACE? (mapping_fd_, graph_fd_)

Honestly, I haven't really thought about what those file descriptors are used for and when they are populated. But yeah, we should probably move those member variables over when creating the new object.

aamijar · 2026-08-12T04:07:19Z


 private:
+  template <typename, typename, ann_dataset_view>
+  friend struct index;


We need this because one index specialization is cannot access the private variables of another. This is the functionality that we need in our new constructor.

aamijar · 2026-08-12T04:08:25Z

- * This returns a new padded index because standard and padded MG indexes have different C++ types.
+ * This moves each rank-local CAGRA graph into the returned padded MG index.
 */
-auto attach_dataset(


attach_dataset has been deleted in favor update_dataset

aamijar · 2026-08-12T04:08:59Z

 /**
 * @brief Update an existing padded MG CAGRA index with a padded dataset of the same layout.
 */
-void update_device_dataset_same_layout(


update_device_dataset_same_layout has been deleted in favor of using the member function index.update_dataset().

aamijar · 2026-08-12T04:11:13Z

+          typename IdxT,
+          ann_dataset_view SrcDatasetViewT,
+          ann_dataset_view DstDatasetViewT>
+auto update_dataset(raft::resources const& res,


This is where the update_dataset implementation lives. Again, it is mostly a wrapper around the newly added constructor.

I think we should move away from state changing user functions directly on the index. How about we add a new overload cagra::update_dataset for which you pass an index by reference (if the dataset type is the same)? Currently the distinction seems unclean -- if the dataset is the same type, call the mutating API, otherwise call the free function.
It looks like currently we don't even have a way for the user to cleanly query the type of dataset on the index, right? So they would be guessing which update_dataset to call.

So you are saying delete the member function and make it a free function instead right?

I agree that seems cleaner, in fact the MG index types don't have any member functions either.

I would argue that technically even with the same dataset type, you are sort of modifying the index itself if you modify the dataset. We could maintain only a single free API that takes an rvalue (the current free API that you already have) of the index and returns a new index even if the dataset type is the same. That would make the behavior more consistent and predictable.

Okay that sounds fine, one question is we are adding a bunch of declarations of update_dataset in cagra.hpp. We are already at 16. From what I understand there isn't a way to get around this because we aren't allowing templates in headers. And the main benefit of that architecture is so that we as cuvs are shipping all the compiled definitions as part of the binary? So the user doesn't need to compile themselves?

Templates could very well have been present. By avoiding them we just restrict what types are allowed at compile time.

We could pre-compile as part of the binary even with templates. But then the user would get a run-time error if they do f<doube>() and double is not suppported.

aamijar · 2026-08-12T04:13:08Z

+  {                                                                                           \
+    return cuvs::neighbors::cagra::update_dataset<T, IdxT, SrcDatasetViewT, DstDatasetViewT>( \
+      res, std::move(cagra_index), dataset);                                                  \
+  }


I am not sure whether this is the correct file where the concrete definitions should live. But it seemed repetitive to create another file like cagra_update_dataset_inst.cu.in.

aamijar · 2026-08-12T04:16:13Z

-    // call cagra::index::update_device_dataset_same_layout on it to update the ann_index to point
-    // to the
-    // new dataset
+    // call cagra::index::update_dataset on it to update the ann_index to point to the new dataset


We need to rework tiered_index IMO. I've added my thoughts here #2434. We will properly get rid of the convert_standard_to_padded_index from cagra.hpp that is only used by tiered index in a follow up.

aamijar · 2026-08-12T04:36:55Z

+auto update_dataset(raft::resources const& res,
+                    index<half, uint32_t, host_standard_dataset_view<half, int64_t>>&& cagra_index,
+                    device_standard_dataset_view<half, int64_t> dataset)
+  -> index<half, uint32_t, device_standard_dataset_view<half, int64_t>>;


The four instantiations for host_standard_dataset_view -> device_standard_dataset_view are to support the MG build implementation.

However, this use case is puzzling to me. Would a user ever need to do this? And why are we doing it in the MG code? We should be doing host_standard_dataset_view -> device_padded_dataset_view if anything.

Another question is why are we doing copying from H2D within the MG code anyway? Why not move that responsibility to the user so that they have ownership of the device dataset?

tarang-jain · 2026-08-13T17:27:52Z

+          typename IdxT,
+          ann_dataset_view SrcDatasetViewT,
+          ann_dataset_view DstDatasetViewT>
+auto update_dataset(raft::resources const& res,


I think we should move away from state changing user functions directly on the index. How about we add a new overload cagra::update_dataset for which you pass an index by reference (if the dataset type is the same)? Currently the distinction seems unclean -- if the dataset is the same type, call the mutating API, otherwise call the free function.
It looks like currently we don't even have a way for the user to cleanly query the type of dataset on the index, right? So they would be guessing which update_dataset to call.

tarang-jain · 2026-08-13T17:32:29Z

+      graph_degree_(other.graph_degree_)
+  {
+    update_dataset(res, dataset);
+  }


Arent we also supposed to move the file descriptors for ACE? (mapping_fd_, graph_fd_)

update_dataset

491ebc0

github-project-automation Bot added this to Unstructured Data Processing Aug 8, 2026

aamijar self-assigned this Aug 8, 2026

aamijar added non-breaking Introduces a non-breaking change improvement Improves an existing functionality labels Aug 8, 2026

aamijar and others added 3 commits August 7, 2026 18:37

Merge branch 'main' into update_dataset

596dd19

remove attach_dataset

eb34e5d

C layer remove attach_dataset

75c3248

aamijar commented Aug 12, 2026

View reviewed changes

aamijar marked this pull request as ready for review August 12, 2026 04:44

aamijar requested review from a team as code owners August 12, 2026 04:44

Merge branch 'main' into update_dataset

3261373

tarang-jain requested changes Aug 13, 2026

View reviewed changes

Uh oh!

Conversation

aamijar commented Aug 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

I've annotated this PR below to make it easier to review.

Uh oh!

copy-pr-bot Bot commented Aug 8, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aamijar Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

aamijar commented Aug 8, 2026 •

edited

Loading

aamijar Aug 12, 2026 •

edited

Loading