Skip to content

update_dataset CAGRA - #2427

Open
aamijar wants to merge 5 commits into
NVIDIA:mainfrom
aamijar:update_dataset
Open

update_dataset CAGRA#2427
aamijar wants to merge 5 commits into
NVIDIA:mainfrom
aamijar:update_dataset

Conversation

@aamijar

@aamijar aamijar commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

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

copy-pr-bot Bot commented Aug 8, 2026

Copy link
Copy Markdown

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 aamijar self-assigned this Aug 8, 2026
@aamijar aamijar added non-breaking Introduces a non-breaking change improvement Improves an existing functionality labels Aug 8, 2026
#include <numeric>
#include <optional>
#include <string>
#include <type_traits>

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cleaning up some unused headers and reordering.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removing this unnecessary forward declaration.

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)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding explicit keyword to avoid unintentional implicit conversions.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).


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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

* 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(

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

attach_dataset has been deleted in favor update_dataset

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

// 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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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>>;

@aamijar aamijar Aug 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

@aamijar
aamijar marked this pull request as ready for review August 12, 2026 04:44
@aamijar
aamijar requested review from a team as code owners August 12, 2026 04:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

improvement Improves an existing functionality non-breaking Introduces a non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FEA][C++ API] Fix update_dataset at C++ API layer

1 participant