[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):

Intermediate Class Structure (Classes + Inheritance):

dataset and dataset_view are now 2 separate parent classes.

Intermediate Class Structure (ContainerType Tags + Composition + Variants):

New Class Structure (ContainerType Tags + Composition):

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

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.

cuvs_cagra_c_api_lifetime_holder is a separate heap object from cagra::index. It is heap-allocated in cagra.cpp with new cuvs_cagra_c_api_lifetime_holder<...>. The C API keeps a raw pointer to it in cuvsCagraIndex.cuvs_cagra_c_api_lifetime_holder It is not embedded in the index, which is why the C layer needs that second field to delete the holder on destroy.

Heap-allocated bundle for the C API: owns cagra::index and any co-owned device storage (VPQ, padded dataset copy, merge/de-serialize/extend buffers) when the index is not standalone. cuvsCagraIndex.c_api_lifetime_owner points at this. Used for merge, build, deserialize, from_args, extend.

The holder moves the owning device_padded_dataset (as unique_ptr<dataset<>> in padded_dataset_owner) to the heap, and cuvsCagraIndex.merged_owner points at the holder. Destroying the C index later destroys the holder first, so the dataset outlives the index’s use of the view, or the ordering is set up so the view is not used after free.

In cuvsCagraIndexFromArgs in cagra.cpp (C API) where callers are things like the Python cagra.from_graph (via Cython) and the Java CagraIndexImpl, and any C code that uses that function:

The flow is: caller → cuvsCagraIndexFromArgs → _from_args, which writes into the cuvsCagraIndex struct the user passed

The holder is not returned as a separate C return value. It is allocated on the heap and its address is stored in output_index->merged_owner, and output_index->addr points at the index inside that holder (or at a freestanding index when merged_owner == 0).

So when _from_args returns, the user’s cuvsCagraIndex already holds the pointers that describe where everything lives.

The unique_ptr to the copy of the dataset from make_padded_dataset is not local to _from_args—it is a member of the holder, which is on the heap and stays alive.

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.

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

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

…A level

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

…tion between make host/device padded dataset in factory

… of dataset + create build_result struct which returns both index and vpq_dataset to prevent automatic out of scope destruction of dataset for vpq case

…rt for cases where we DO need to own the dataset (in order to keep view alive for index). All cases where we build() from dataset already on device --> we don't need to own. Merge + All cases when data is on host --> we DO need to own the device copy we create. This includes within ACE build and C API build from host and from_args with host dataset

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

…d to declare the output index type passed into deserialize(). User allocates memory outside this call and owns this memory since index doesn't own dataset anymore. Since memory is passed into deserialize, deserialize will already know the type of index it should deserialize at compile time because the user specifies it. This is the new API contract and is done in iface.hpp. Adding additional index types for deserialize will just require an additional dispatch branch in the iface in the future

…hit 2 argument index not recognized error. This is because index used to be untyped with 2 args but now we added DatasetViewT template param on index which made index 3 args so FAISS call sites for index which were still 2 args were no longer recognized.

HowardHuang1 · 2026-06-11T03:55:00Z

/ok to test 4258dc5

…_view for ACE build The new cagra::build() API requires dataset_view types (host_padded_dataset_view or device_padded_dataset_view) rather than raw raft::mdspan. This fixes the following build errors: - "no instance of function template cagra::attach_device_dataset_on_host_index matches" (caused by cagra::build resolving to mg_index overload when passed raw mdspan with uint32_t extents, since index_params implicitly converts to mg_index_params) - "no operator = matches these operands" for *dataset_ = std::move(padded->data_) (caused by uint32_t vs int64_t mdarray extents mismatch) Fix: use int64_t extents for dataset_view_host so make_device_padded_dataset produces data_ with matching int64_t extents, and explicitly create an owning host_padded_dataset (or zero-copy view when stride already matches) for the ACE cagra::build calls.

Wrap device dataset with make_device_padded_dataset_view before passing to cuvs::neighbors::cagra::build(), and use make_device_padded_dataset with as_dataset_view() for the host-data path, satisfying the is_device_dataset_view constraint required by the new build() API.

…isk-mode hnsw.build Three related bugs caused test_cagra_ace.py and test_hnsw_ace.py failures: 1. Use-after-free in cagra C API (in-memory ACE, ~1% recall): convert_host_to_device_index stores a VIEW of host_idx.graph_ (no ownership transfer). When host_idx goes out of scope, device_idx.graph_view_ becomes dangling. Fix: copy device graph to host then back to a new owned device copy before host_idx is destroyed. 2. Missing graph_fd/mapping_fd transfer in cagra C API (disk-mode ACE): After convert_host_to_device_index, only dataset_fd was stolen from host_idx. The graph_fd and mapping_fd were never transferred, so hnsw::from_cagra's disk-path check (dataset_fd && graph_fd) always fell through to the GPU path which failed with "No dataset provided". Fix: add steal_graph_fd()/steal_mapping_fd() to the index class and transfer all three FDs in the C API disk-mode block. 3. hnsw::build with use_disk=True never wrote hnsw_index.bin: The ACE host index had dataset/graph/mapping FDs, but the code created an in-memory device padded dataset and called attach_device_dataset_on_host_index, discarding the FDs. from_cagra never saw a disk-backed index and never called serialize_to_hnswlib_from_disk. Fix: detect disk mode (dataset_fd present), transfer all FDs to a fresh device_padded_index, then call from_cagra which takes the disk path and serializes the HNSW index to build_dir/hnsw_index.bin.

The CuvsCagra<data_t>::train() hunk had wrong counts in its @@ header: @@ -204,14 +207,22 @@ (declared 14 before, 22 after) @@ -204,15 +207,21 @@ (actual 15 before, 21 after) The off-by-one counts in opposite directions caused `git apply` to report "corrupt patch at line 301", silently skipping the entire 26.06 diff. This left GpuResources.h with the old rmm/mr/device_memory_resource.hpp include that was removed in RMM 26.06, causing CI build failures. Also removed a spurious blank line between the train() and search() hunks that caused a "patch fragment without header" error.

HowardHuang1 · 2026-06-11T21:49:43Z

/ok to test 9c024c6

…ce path. Host_padded_index cannot have attach_dataset_on_build=true since we can't search on a host dataset regardless. ACE non-disk, ACE sub-builds, and merge no longer need attach_dataset_on_build in the new design

…taset_view The device path in BinaryCuvsCagra::train() was passing a raw mdspan to cagra::build() instead of a device_padded_dataset_view, causing: error: no instance of overloaded function "cuvs::neighbors::cagra::build" matches the argument list Fix by wrapping the device matrix view with make_device_padded_dataset_view before passing to build(), matching the pattern used in CuvsCagra<data_t>::train(). Also corrects the downstream hunk offsets (+183/+225/+297) in the BinaryCuvsCagra search() and reset() sections that were shifted by the 3-line addition.

update_graph(device_matrix_view) stores only a VIEW with no ownership transfer. When convert_host_to_device_index is called with a temporary host_idx (e.g. inside iface.hpp for the MG host-data build path), the returned device index ends up with a dangling graph_view_ once host_idx is destroyed, causing a segfault on first access. Fix: copy the graph device→host→device inside convert_host_to_device_index so the returned index always owns its graph memory. Remove the now-redundant workaround that was doing the same extra copy in cagra.cpp.

HowardHuang1 · 2026-06-12T01:26:34Z

/ok to test e795819

…ed update_dataset() supporting host_matrix and device_matrix views were removed

HowardHuang1 · 2026-06-15T00:38:21Z

/ok to test 26b87e5

…iew. All call sites updated to use Dataset API instead

…The main library already uses C++20

…ncepts. The main library already uses C++20" This reverts commit 48fc8d5.

HowardHuang1 · 2026-06-16T17:31:06Z

/ok to test 217e5bd

… to use the new Dataset API

HowardHuang1 · 2026-06-16T19:42:28Z

/ok to test 8001bf4

HowardHuang1 · 2026-06-25T03:36:47Z

/ok to test 3803e00

HowardHuang1 · 2026-06-25T05:11:53Z

/ok to test c31ea7c

HowardHuang1 added 4 commits February 23, 2026 09:49

get build working

d78f459

add dataset compression test and basic constructor types test

4febf8b

add padded_dataset class along with test cases

b403473

add support for new padded_dataset classes all the way up to the CAGR…

8d6833a

…A level

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

Merge branch 'main' into HH-Dataset-API

5447a4c

fix style

17ab09d

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

build() now only takes views and not unique ptrs + get rid of distinc…

fb556c9

…tion between make host/device padded dataset in factory

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

HowardHuang1 added 4 commits March 2, 2026 15:27

clean up old overloads of build & index functions that take ownership…

37d28dc

… of dataset + create build_result struct which returns both index and vpq_dataset to prevent automatic out of scope destruction of dataset for vpq case

fix failing mg tests that do build -> serialize -> deserialize -> search

26b46a2

Merge remote-tracking branch 'upstream' into HH-Dataset-API

a38fb18

HowardHuang1 added 3 commits June 10, 2026 19:09

add fixes at call sites once variants were removed completely

828c371

HowardHuang1 added 5 commits June 11, 2026 10:41

Merge remote-tracking branch 'upstream' into HH-Dataset-API

0dfeffe

HowardHuang1 added 4 commits June 11, 2026 16:21

minor fix for update_graph overload confusion

e795819

fix FAISS to call update_dataset() with padded dataset after deprecat…

26b87e5

…ed update_dataset() supporting host_matrix and device_matrix views were removed

HowardHuang1 added 6 commits June 14, 2026 17:39

remove deprecated update_dataset() overload taking in device_matrix_v…

2804440

…iew. All call sites updated to use Dataset API instead

update examples/cpp/CMakeLists.txt to use C++20 to support concepts. …

48fc8d5

…The main library already uses C++20

Revert "update examples/cpp/CMakeLists.txt to use C++20 to support co…

ff8cc82

…ncepts. The main library already uses C++20" This reverts commit 48fc8d5.

move FAISS changes for new Dataset API into separate 26.08 patch

4d7d0a6

use raft::copy_matrix instead of cudaMemcpy2DAsync

bb275f1

Merge remote-tracking branch 'upstream' into HH-Dataset-API

217e5bd

fix two examples that still use the old cagra API. Migrated them over…

8001bf4

… to use the new Dataset API

Merge remote-tracking branch 'upstream' into HH-Dataset-API

3803e00

fix merge conflict. Recover missing line

c31ea7c

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

Intermediate Class Structure (Classes + Inheritance):

Intermediate Class Structure (ContainerType Tags + Composition + Variants):

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:

3 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!

HowardHuang1 commented Jun 11, 2026

Uh oh!

HowardHuang1 commented Jun 11, 2026

Uh oh!

HowardHuang1 commented Jun 12, 2026

Uh oh!

HowardHuang1 commented Jun 15, 2026

Uh oh!

HowardHuang1 commented Jun 16, 2026

Uh oh!

HowardHuang1 commented Jun 16, 2026

Uh oh!

HowardHuang1 commented Jun 25, 2026

Uh oh!

HowardHuang1 commented Jun 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

HowardHuang1 commented Feb 24, 2026 •

edited

Loading

HowardHuang1 commented Mar 4, 2026 •

edited

Loading