[REVIEW] New Dataset API Clarifying Ownership by HowardHuang1 · Pull Request #1846 · NVIDIA/cuvs

HowardHuang1 · 2026-02-24T23:21:21Z

Overview

Addressing #1574 and #1571.

Replaced strided_dataset with padded_dataset class. Added support all the way up to CAGRA code.

Old class structure (Classes + Inheritance):

New Class Structure (ContainerType Tags + Composition):

Inheritance is removed entirely and all dataset types are on the same level of the inheritance tree.

3 Levels:

Storage

Container tags

Public dataset/view aliases

Ownership

The index and cagra::build / cagra::index do not own raw vector storage, they only take views.

callers (or the C merged holder) must keep backing memory alive for as long as the index is used.

The old code had a type-erased std::unique_ptr<dataset_view<...>>, i.e. non-owning view handles. The new code uses templates on the index type which determines the type of dataset_view the index holds.

ACE v.s. non-ACE paths on Host

ACE path copies datasets that can't entirely fit in CPU memory in chunks onto GPU memory by calling make_padded_dataset. This is 1x memory on CPU and 1x memory on GPU.

Return types:

Used mainly to maintain lifetime of dataset.

cuvs_cagra_c_api_lifetime_holder

unique_ptr<vpq_dataset> vpq_owner
unique_ptr padded_dataset_owner
raft::device_matrix dataset
cagra::index idx
It is a single C++ struct in cagra.cpp that groups the real cagra::index with any extra heap-owned things the C API had to create so the index’s non-owning views stay valid.

Miscellaneous: Extend Serialize Deserialize

Will fill in later

Factories:

make_device_padded_dataset_view
make_host_padded_dataset_view
make_device_padded_dataset
make_host_padded_dataset
make_vpq_dataset
- in pq.hpp and pq.cu
make_merged_dataset
- in cagra.hpp

Places where make_padded_dataset/view are called internally (not by user):

Host non-ACE path

cpp/src/neighbors/cagra_build_inst.cu.in
cagra_from_host_padded in cpp/src/neighbors/iface/iface.hpp
c/src/neighbors/cagra.cpp

Tiered CAGRA

update_cagra_ann_dataset_for_stride
build_upstream_ann

Ownership in Downstream Functions:

build() takes dataset_view only.
Downstream functions search / serialize / deserialize / merge only take views.

Improvements:

build() function previously only supported device dataset inputs. It now supports host dataset inputs.

Breaking Changes for Dataset API:

The following functions are removed since index no longer owns the dataset, index only takes views:

Removed all owning dataset based builds. Build only takes views.
Removed all update_dataset() overloads that take owning dataset. Update_dataset() only takes views.
Removed old functions that took mdspan or derivatives of mdspan.

4 cases where index previously owned dataset [all deprecated paths]:

2 edge case build() paths when attach_dataset_on_build == true and a successful dense attach:

Non-ACE / typical padded attach: rows live under index_owning_dataset_storage_ (type-erased owning wrapper, commonly device_padded_dataset).
ACE in-memory device_matrix attach: rows live under index_owning_dataset_storage_ (optional raw device_matrix).

Compression Param:

implicit vpq dataset creation within build() when compression==True
this forces index to own new vpq dataset which violates our new contract that we want index to never own dataset.

Merge:

MERGE path: merge() internally creates merged_dataset on a deprecated internal merged dataset creation path. Here, index takes ownership of merged_dataset by storing it in index_owning_dataset_storage_ .

These paths have since been removed.

Attach Dataset

Previously, in the old code, ACE attach_dataset_on_build called make_device_padded_dataset on host dataset which did a H2D copy in order to attach dataset to final index.
cpp/src/neighbors/detail/cagra/cagra_build.cuh
This has since been removed. attach_dataset_on_build is disabled for host build paths. This avoids a H2D copy.

Compressed Dataset

Removed old code that did compressed dataset creation within build. This should only happen in the factory.

Merged Dataset

Removed implicit memory allocation within merge(), memory allocation now delegated to make_merged_dataset() factory. Removed index ownership within merge.

Deserialize

Removed index ownership of dataset during deserialization. Now users are expected to create/declare the dataset type to be deserialized and then pass it as a reference to the deserialize() function which will then populate this dataset and return it to the caller.

Helpers

cagra_required_row_width
matrix_actual_row_width
matrix_row_width_matches_cagra_required
convert_dataset_view_to_padded_for_graph_build
convert_host_to_device_index
attach_device_dataset_on_host_index

How to attach a compressed dataset onto an uncompressed index?

construct a new compressed index
copy over the graph and other params from the uncompressed index
delete the old uncompressed index
attach the vpq dataset onto the compressed index

How to attach a searchable device dataset onto an index built with host build?

Convert host index to device index with helper function convert_host_to_device_index
a. Utilizes map of host dataset type to device dataset type counterpart

TODOs:

Bring back Host functions [DONE]
Mark any old functions that are no longer used as [DONE]
Use templates wherever possible. Shift towards composition rather than inheritance [DONE]

Recent Updates:

build_ace() and build() functions merged on Public API surface
removed build_result(), ace_build_result(), and merge_result()
deprecated internal vpq dataset creation inside build() when index_params::compression == true --> moved to make_vpq_dataset() factory
deprecated internal merged dataset creation inside merge() --> moved to make_merged_dataset() factory. For backwards compatibility, have index take ownership of deprecated internal merged dataset creation path ONLY.
build() and build_ace() both had a attach_dataset_on_build which requires ownership of dataset. Ownership is given to index temporarily. This will later be deprecated. Users will be expected to pass padded dataset on device and call search() directly. Attach_dataset_on_build will no longer be supported for host builds.
added host versions of dataset API
templated build() and downstream functions to work on host datasets
added index template type conversion helpers

Future PRs:

PR#2: Add Support for Compressed Datasets

pq_dataset
bbq_dataset
rabitq_dataset
sq_dataset

PR#3: Migrate Rest of Algorithms to use Dataset API

HNSW
IVF
Vamana
Scann
Brute Force

copy-pr-bot · 2026-02-24T23:21:25Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

aamijar · 2026-02-25T01:51:35Z

/ok to test 5447a4c

aamijar · 2026-02-25T02:12:45Z

/ok to test 17ab09d

achirkin · 2026-02-25T06:06:02Z

NB: I updated the label to breaking, since the description implies removal of a publicly visible class strided_dataset

cjnolet · 2026-02-25T22:22:05Z

Does the dataset(_view) type bring anything on top of mdarray/mdspan in that case?

@achirkin The problem w/ using mdspan/mdarray for this is that it's not carrying along the proper information to either the algorithms nor the user (which is why we created this specialized class for this in the first place!).

Two immediate reasons why this API is necessary:

The user should not have to know that they need to pad a dataset in order to use cagra without the additional copy. They should not need to know how any of these algorithms work internally. They should, however, need to know that CAGRA expects a padded dataset, and they should have an API to construct one so that they can own the dataset class and not have cagra creating one under the hood.
APIs, especially the graph-based APIs, should be able to accept as inputs data which has been quantized using a metod like PQ, which carries with additional information. In the case of PQ, the codebooks are needed to compute the distances. This again decouples the quantization from the algorithm (CAGRA-Q does not need to do its own quantization. It should just accept the quantized vectors). We're being asked for the same behavior with Vamana.

This new API solves both of these problems while leaving the control over the memory ownership entirely in the user's hands. We've discussed this for a long time. We've known this is needed for a long time. it's time to prioritize this and get it done. I agree that an anstract class might make more sense, but ultimately we should not be moving any owneship over to the algorithm (the user should maintain ownership over the class and underlying memory the entire time).

HowardHuang1 · 2026-03-04T23:24:04Z

The doc that outlines some of the API design choices can be found in slack. Let me know if there are any parts of the design that can be altered to better suit our users' needs.

The following files are test case files I've added and can be ignored for now. They will be removed before the final merge with upstream repo:

cagra_build_view_only.cu
cagra_padded_dataset.cu
cagra_vpq_build_result.cu
dataset_compression.cu
dataset_types.cu

cjnolet · 2026-03-25T14:54:12Z

/ok to test 70b6b58

achirkin

Sorry for being so late for the second round of reviews! I have a few suggestions for the dataset type hierarchy.

…ets: additional dataset and final extended_dataset to be populated. Extend() function no longer internally does H2D copy. Now, extend() extends graph only. Users are expected to concat original dataset + additional dataset themselves into 1 single extended dataset which they then pass in to extend() function along with a variable called new_start_row indicating at which index the additional dataset starts. Extend() now only accepts device_padded_dataset inputs since it internally calls search kernels that fail on any other dataset type. Users must do H2D copy themselves with the concatenated dataset prior to passing it in to extend()

HowardHuang1 · 2026-07-30T01:15:41Z

/ok to test c8df0d4

…-API # Conflicts: # c/tests/neighbors/ann_cagra_c.cu # cpp/include/cuvs/neighbors/cagra.hpp # cpp/src/neighbors/cagra_search_inst.cu.in # cpp/tests/neighbors/ann_cagra.cuh

Restore java serialize behavior

HowardHuang1 · 2026-07-30T04:19:47Z

/ok to test 5581ffd

… API with strided_dataset and build functions that don't take dataset_view. Migrated call sites to new Dataset API

HowardHuang1 · 2026-07-30T05:58:43Z

/ok to test 8140998

…ose call sites over to use new Dataset API

HowardHuang1 · 2026-07-30T06:39:36Z

/ok to test 62a8da6

cjnolet · 2026-07-29T20:58:13Z

+#pragma once
+
+#include <cuvs/core/c_api.h>
+#include <cuvs/neighbors/common.h>


This should not be included in here. We shouldn't including things across different namespaces / top level directories.

I removed common.hpp here 02963e1

cjnolet · 2026-07-29T22:06:46Z

        return "Index(type=CAGRA, metric=L2" + (", ".join(attr_str)) + ")"


+cdef class Dataset:


Why is this in CAGRA? This should be a separate file in core/common to mirror the other common vocab types

cjnolet · 2026-07-29T22:07:43Z

+                self.dataset.dtype.lanes)
+
+
+cdef class PaddedDataset(Dataset):


Why? This is just an empty class. what's the purpose of having a separate type at all if it's not going to be any different from the other types other than the name? This doesn't make sense to me.

cjnolet · 2026-07-29T22:08:15Z

+
+
+cdef class PaddedDatasetView:
+    def __cinit__(self):


Same with these. If we don't care baout the type, don't bother making a separate type. We gain nothing but bloat and confusion by doing this.

cjnolet · 2026-07-29T22:09:17Z

+
+
 @auto_sync_resources
 def build(IndexParams index_params, dataset, resources=None):


In Python we don't care about views vs owning, so please don't make this distinction here... The user shouldn't need to create a "view" to use CAGRA. We should be able to make that behind the scenes if needed.

cjnolet

Had a bunch of pending review items.

cjnolet · 2026-07-30T15:07:45Z

+typedef cuvsDatasetView* cuvsDatasetView_t;
+
+/**
+ * @brief Create an empty owning dataset handle.


We can improve the docs here slightly to also mention that the underlying data is not allocated, only the handle.

aamijar · 2026-07-30T17:08:36Z

/ok to test 378ac4e

aamijar · 2026-07-30T17:24:35Z

/ok to test 02963e1

cjnolet · 2026-07-30T17:44:34Z

+ * @param[out] layout   resolved dataset layout
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsCagraGetDatasetMemTypeAndLayout(DLManagedTensor* dataset,


@tarang-jain suggested we temove this and I agree. We don't need to know what type the underlying cagra index holds. Thast should be handled. We should howver, have accessors for the dataset mem type and layout, but that should be in dataset.h.

cjnolet · 2026-07-30T17:45:19Z

 * @param[in,out] index cuvsCagraIndex_t CAGRA index
 * @return cuvsError_t
 */
 CUVS_EXPORT cuvsError_t cuvsCagraExtend(cuvsResources_t res,


Need to make sure we have tests for this- especially tests that use the extended_dataset. @HowardHuang1 to make sure we have C tests for this.

cjnolet · 2026-07-30T17:45:30Z

+ * @param[in] filename the file name for saving the graph
+ * @param[in] index CAGRA index
+ */
+CUVS_EXPORT cuvsError_t cuvsCagraSerializeGraph(cuvsResources_t res,


Need to mke sure we have tests.

cjnolet · 2026-07-30T18:06:49Z

/ok to test 09b3334

cjnolet · 2026-07-30T18:13:02Z

+  ASSERT_EQ(cuvsCagraBuild(res, build_params, padded_dataset_owner, owner_built_index), CUVS_SUCCESS);
+  ASSERT_EQ(cuvsCagraIndexDestroy(owner_built_index), CUVS_SUCCESS);
+  ASSERT_EQ(cuvsCagraUpdateDataset(res, padded_dataset_owner, index), CUVS_SUCCESS);
+  ASSERT_EQ(cuvsCagraUpdateDataset(res, padded_dataset_owner, index), CUVS_SUCCESS);


What is calling this twice demonstrating? Is that intentional? If so, we should add a comment here as to why, and I'd consider pulling this out into a separate test so that it's immediatley obvious to future eyes why this is called twice.

The reason I flag this specifically is because it's bound to get removed at some point when future eyes come into this file and assume it's a bug in the test.

cjnolet · 2026-07-30T18:19:14Z

+  host_tensor.dl_tensor.shape              = dataset_shape;
+
+  cuvsDataset_t host_view = nullptr;
+  ASSERT_EQ(cuvsDatasetMakePaddedView(res, &host_tensor, &host_view), CUVS_SUCCESS);


We assert success everywhere but we don't seem to assert any failure cases. We should always test what we want to work.... that means testing for the failures too.

cjnolet · 2026-07-30T18:21:02Z

+
+  auto serialized_path =
+    "/tmp/cuvs-cagra-host-padded-" + std::to_string(getpid()) + ".bin";
+  ASSERT_EQ(cuvsCagraSerializeGraphAndDataset(res, serialized_path.c_str(), index), CUVS_SUCCESS);


This test is testing too much. This is supposed to be testing update. We should have separate tests for serialization and deserialization. The reason for having separate test cases is so that when they fail, it's immediately obvious why "Oh the update broke" vs "oh the update test broke but it seems to be because something happened in the serialization part of the test". The former is immediately obvious, the latter requires a much closer look

Just making sure I'm clear here- it's okay to have tests that call multiple functions, but that shouldn't be at the expense of having a smaller more targeted test in addition to the larger more integration-centric tests.

cjnolet · 2026-07-30T18:23:22Z

+  cuvsCagraIndex_t loaded_index;
+  ASSERT_EQ(cuvsCagraIndexCreate(&loaded_index), CUVS_SUCCESS);
+  cuvsDataset_t loaded_dataset = nullptr;
+  ASSERT_EQ(cuvsCagraDeserializeGraphAndDataset(


Is the purpose of adding this to the test so that we can create a host padded dataset? Is that the only way to construct this type of dataset? That seems to indicate a bad smell in the usage of the APIs. I think we should be able to specify "I want to create a host padded dataset" if needed.

HowardHuang1 requested review from a team as code owners February 24, 2026 23:21

github-project-automation Bot added this to Unstructured Data Processing Feb 24, 2026

aamijar added non-breaking Introduces a non-breaking change feature request New feature or request labels Feb 25, 2026

aamijar assigned HowardHuang1 Feb 25, 2026

aamijar moved this to In Progress in Unstructured Data Processing Feb 25, 2026

achirkin added breaking Introduces a breaking change and removed non-breaking Introduces a non-breaking change labels Feb 25, 2026

achirkin reviewed Feb 25, 2026

View reviewed changes

Comment thread cpp/include/cuvs/neighbors/cagra.hpp

aamijar reviewed Feb 25, 2026

View reviewed changes

Comment thread cpp/CMakeLists.txt Outdated

seunghwak mentioned this pull request Feb 27, 2026

[WIP] Clarify dataset ownership and allocation semantics #1738

Closed

HowardHuang1 requested a review from a team as a code owner February 28, 2026 03:34

cjnolet mentioned this pull request Mar 10, 2026

[FEA] Implement binary dataset and accept in CAGRA #1906

Open

mohanprasand-nuvai mentioned this pull request Mar 23, 2026

Binary dataset support in CAGRA (upstream #1906) Nuvai/cuvs#5

Closed

5 tasks

HowardHuang1 changed the title ~~[WIP] New Dataset API Clarifying Ownership~~ [REVIEW] New Dataset API Clarifying Ownership Mar 25, 2026

HowardHuang1 changed the base branch from main to release/26.04 March 25, 2026 01:03

HowardHuang1 requested review from a team as code owners March 25, 2026 01:22

HowardHuang1 requested a review from bdice March 25, 2026 01:22

tfeher requested changes Apr 1, 2026

View reviewed changes

Comment thread cpp/include/cuvs/neighbors/cagra.hpp Outdated

Comment thread cpp/include/cuvs/neighbors/cagra.hpp Outdated

tfeher changed the base branch from release/26.04 to main April 1, 2026 23:31

achirkin requested changes Apr 8, 2026

View reviewed changes

divyegala and others added 3 commits July 30, 2026 00:20

add datasets back

fe7a121

Merge remote-tracking branch 'origin/java-serialize' into java-serialize

f046dd2

HowardHuang1 added 2 commits July 29, 2026 21:13

Merge remote-tracking branch 'upstream/release/26.08' into HH-Dataset…

e916441

…-API # Conflicts: # c/tests/neighbors/ann_cagra_c.cu # cpp/include/cuvs/neighbors/cagra.hpp # cpp/src/neighbors/cagra_search_inst.cu.in # cpp/tests/neighbors/ann_cagra.cuh

Merge pull request #3 from divyegala/java-serialize

5581ffd

Restore java serialize behavior

Fix merge conflicts. Upstream introduced new test cases that used old…

8140998

… API with strided_dataset and build functions that don't take dataset_view. Migrated call sites to new Dataset API

More test case call site fixes to resolve merge conflicts and port th…

62a8da6

…ose call sites over to use new Dataset API

cjnolet requested changes Jul 30, 2026

View reviewed changes

cjnolet reviewed Jul 30, 2026

View reviewed changes

Comment thread c/include/cuvs/core/dataset.h Outdated

aamijar added 2 commits July 30, 2026 16:15

remove explicit view type in C API

06b02ad

bool is_owning

378ac4e

remove unused header in dataset.h

02963e1

dantegd approved these changes Jul 30, 2026

View reviewed changes

cjnolet reviewed Jul 30, 2026

View reviewed changes

Increase binary limits to pass CI

09b3334

cjnolet reviewed Jul 30, 2026

View reviewed changes

		return "Index(type=CAGRA, metric=L2" + (", ".join(attr_str)) + ")"


		cdef class Dataset:



		@auto_sync_resources
		def build(IndexParams index_params, dataset, resources=None):

Uh oh!

Conversation

HowardHuang1 commented Feb 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Old class structure (Classes + Inheritance):

New Class Structure (ContainerType Tags + Composition):

Ownership

ACE v.s. non-ACE paths on Host

Return types:

Miscellaneous: Extend Serialize Deserialize

Factories:

Places where make_padded_dataset/view are called internally (not by user):

Ownership in Downstream Functions:

Improvements:

Breaking Changes for Dataset API:

4 cases where index previously owned dataset [all deprecated paths]:

Attach Dataset

Compressed Dataset

Merged Dataset

Deserialize

Helpers

How to attach a compressed dataset onto an uncompressed index?

How to attach a searchable device dataset onto an index built with host build?

TODOs:

Recent Updates:

Future PRs:

Uh oh!

copy-pr-bot Bot commented Feb 24, 2026

Uh oh!

aamijar commented Feb 25, 2026

Uh oh!

aamijar commented Feb 25, 2026

Uh oh!

achirkin commented Feb 25, 2026

Uh oh!

Uh oh!

cjnolet commented Feb 25, 2026

Uh oh!

Uh oh!

HowardHuang1 commented Mar 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cjnolet commented Mar 25, 2026

Uh oh!

Uh oh!

Uh oh!

achirkin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

HowardHuang1 commented Jul 30, 2026

Uh oh!

HowardHuang1 commented Jul 30, 2026

Uh oh!

HowardHuang1 commented Jul 30, 2026

Uh oh!

HowardHuang1 commented Jul 30, 2026

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

cjnolet left a comment

Choose a reason for hiding this comment

HowardHuang1 commented Feb 24, 2026 •

edited

Loading

HowardHuang1 commented Mar 4, 2026 •

edited

Loading