Migrate CAGRA merge() to caller-concatenated dataset + offsets contract by HowardHuang1 · Pull Request #2507 · NVIDIA/cuvs

HowardHuang1 · 2026-08-26T04:26:56Z

Addresses issue #2401.

merge() previously allocated the merged dataset buffer and copied each input index's rows into it internally, hiding memory allocation and copies from the caller. This migrates it to the same explicit caller-owned-buffer contract already used by extend(): callers concatenate every input index's dataset (applying any row filter) into a single buffer themselves, and pass per-index offsets marking where each index's rows start. merge() now only builds/merges the graph and binds the returned index to that buffer.

A new merged_dataset_offsets() helper is added for the bitset-filtered case, where the caller can't otherwise derive per-index surviving row counts; unfiltered callers can compute offsets directly (cumulative index sizes) without it.

C++: cagra::merge() and Fastener/rebuild internals updated; new cagra::merged_dataset_offsets().
C API: cuvsCagraMerge/cuvsCagraMergeWithParams updated to match; new cuvsCagraMergedDatasetOffsets.
Java: CagraIndex.merge() public API updated to the same contract (breaking); real in-repo consumer (cuvs-lucene) and tests updated.
Rust, Python, Go: new merge wrappers added (none existed before), matching the C API's contract and modeled on each language's existing extend()/update_dataset() conventions.

Also fixes an unrelated pre-existing build break in cpp/src/core/bloom_filter.cu (cuco::default_filter_policy renamed to cuco::bloom_filter_policy in the pinned cuCollections version), needed to get any build of libcuvs compiling in this environment.

copy-pr-bot · 2026-08-26T04:26:59Z

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.

merge() previously allocated the merged dataset buffer and copied each input index's rows into it internally, hiding memory allocation and copies from the caller. This migrates it to the same explicit caller-owned-buffer contract already used by extend(): callers concatenate every input index's dataset (applying any row filter) into a single buffer themselves, and pass per-index offsets marking where each index's rows start. merge() now only builds/merges the graph and binds the returned index to that buffer. A new merged_dataset_offsets() helper is added for the bitset-filtered case, where the caller can't otherwise derive per-index surviving row counts; unfiltered callers can compute offsets directly (cumulative index sizes) without it. - C++: cagra::merge() and Fastener/rebuild internals updated; new cagra::merged_dataset_offsets(). - C API: cuvsCagraMerge/cuvsCagraMergeWithParams updated to match; new cuvsCagraMergedDatasetOffsets. - Java: CagraIndex.merge() public API updated to the same contract (breaking); real in-repo consumer (cuvs-lucene) and tests updated. - Rust, Python, Go: new merge wrappers added (none existed before), matching the C API's contract and modeled on each language's existing extend()/update_dataset() conventions. Also fixes an unrelated pre-existing build break in cpp/src/core/bloom_filter.cu (cuco::default_filter_policy renamed to cuco::bloom_filter_policy in the pinned cuCollections version), needed to get any build of libcuvs compiling in this environment. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

HowardHuang1 · 2026-08-26T14:28:52Z

/ok to test 87a0075

…taset-to-take-inputs-concat-buffer-and-offset-26_10

HowardHuang1 · 2026-08-26T19:22:32Z

/ok to test 50dd606

…taset-to-take-inputs-concat-buffer-and-offset-26_10

Rust: regenerate cuvs-sys bindings.rs for the offsets-based merge contract, and switch the merge tests from PaddedDataset::new (owning copy, rejects already-aligned sources) to DatasetView::new, mirroring how Index::build picks a view. Go: add MakePaddedDatasetAuto, mirroring BuildIndex's existing padded-vs-standard branch, so merge_test.go can build its merged buffer without needing its own cgo import (which isn't allowed in _test.go files). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…taset-to-take-inputs-concat-buffer-and-offset-26_10

HowardHuang1 · 2026-08-27T22:06:05Z

/ok to test a83394f

…-merge-dataset-to-take-inputs-concat-buffer-and-offset-26_10 # Conflicts: # cpp/src/neighbors/detail/cagra/cagra_merge.cuh # cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu # java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java # java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java # java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java # java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java # java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java # java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java # java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java

…-merge-dataset-to-take-inputs-concat-buffer-and-offset-26_10

HowardHuang1 · 2026-09-14T06:24:57Z

/ok to test 8054382

aamijar

Hi Howard, as discussed offline, let's consider the purpose of this merge PR change.
We already enforce that the user must pre allocate and pass in the merged_dataset storage. The merge() api will populate that storage.
I'm a little confused as to what we are trying to achieve in this PR.

… private detail::build_from_device_matrix()

HowardHuang1 · 2026-09-14T23:42:05Z

/ok to test 7a4845d

…CI can pass

HowardHuang1 · 2026-09-15T17:57:53Z

/ok to test a14a1da

HowardHuang1 · 2026-09-18T00:13:58Z

This PR is more about memory savings. It is also meant to mirror how we altered the extend API. Merge should now be merging graph only.

There are 2 cases: unfiltered and filtered merge.

Unfiltered merge: Peak memory is the same in main and this PR.
Filtered merge: The memory win is the filtered path.

On main:

You allocate GPU storage for surviving rows (final_rows).
Merge allocates a second GPU staging buffer for the unfiltered concat (all rows of every input).
It copies every input dataset into that staging buffer (deleted/filtered rows included).
It applies the bitset and copies survivors into your smaller buffer.
It builds the graph on that smaller buffer.

Peak GPU memory is roughly unfiltered concat + filtered result, plus CSR/index scratch for the bitset. The large staging buffer is merge’s, not yours, and it exists only to throw most of it away.

This PR:

No intermediate staging buffer is created.

Concat is done on CPU by the caller
You allocate only the surviving filtered rows on GPU

Peak GPU memory is just filtered concat. No 2 buffers exist at the same time.

extend()	merge() (this PR)
Caller fills extended_dataset = old \|\| new	Caller fills merged_dataset = index0 \|\| index1 \|\| … (already filtered)
new_start_row marks the split	offsets mark each index’s split
Library does not copy dataset rows	Library does not copy dataset rows
Graph work + update_dataset to bind the view	Graph work + update_dataset to bind the view
Caller owns lifetime	Caller owns lifetime

Hope that helps!

…-merge-dataset-to-take-inputs-concat-buffer-and-offset-26_10

HowardHuang1 · 2026-09-20T05:48:34Z

/ok to test ab5fe77

HowardHuang1 · 2026-09-21T04:53:30Z

/ok to test a8e86d7

cjnolet · 2026-09-21T15:59:26Z

                                                 size_t num_indices,
                                                 cuvsFilter filter,
                                                 cuvsDataset_t merged_dataset,
+                                                 const int64_t* offsets,


This breaks ABI. We can't change C functions (add parameters, etc..).

…-merge-dataset-to-take-inputs-concat-buffer-and-offset-26_10

…rge helpers Recovers the concat/gather logic that used to live inside merge_rebuild before the caller-owned merged_dataset migration, and repackages it as two standalone public helpers callers can use ahead of merge(): concatenate_datasets() for the unfiltered case, and concatenate_and_filter_datasets() for a bitset row_filter. concatenate_and_filter_datasets() does not reintroduce the old merge_rebuild's unfiltered-staging-buffer step -- each index's surviving rows are gathered directly out of that index's own device-resident dataset into the final buffer, so peak device memory stays at the filtered row count only, preserving this PR's core memory-savings goal for the filtered path. Also updates the existing AnnCagraIndexMergeTest/ AnnCagraIndexFilteredMergeTest suites to call the new helpers instead of their own hand-rolled concat/gather code, so they double as end-to-end coverage for both. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…asets() Wraps the two C++ convenience helpers as cuvsCagraConcatenateDatasets and cuvsCagraConcatenateAndFilterDatasets, mirroring the existing cuvsCagraMerge/cuvsCagraMergedDatasetOffsets pattern: dtype dispatch by DLDataType, device-padded/device-standard layout dispatch by input index layout, and output bound as a new owning cuvsDataset_t handle (always device-padded layout, matching the C++ return type). Adds CagraC.ConcatenateDatasetsMergeSearch covering both the unfiltered and bitset-filtered paths end to end (concatenate -> merge -> search), including a MergedDatasetOffsets cross-check. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

dantegd · 2026-09-28T22:12:54Z

-        cuvs::neighbors::cagra::merge(
-          *res_ptr, params_cpp, index_ptrs, view, merge_params, row_filter);
+  using owner_t = cuvs::neighbors::owning_dataset_for_view_t<DatasetViewT>;
+  with_dataset_view<owner_t, DatasetViewT>(merged_dataset, [&](auto const& view) {


Should we validate the merged dataset’s dtype against the input indices before this cast? A float16 buffer can pass the shape/layout checks for float32 indices and then get read as float32, potentially past the allocation. A mismatched dtype test would help here too.

dantegd · 2026-09-28T22:52:47Z

-    // The copy below overwrites columns [0, dim), while the sorter computes L2 over [0, stride).
-    // Zero the remaining padded columns; when stride == dim, there are none to initialize.
-    if (stride > preflight.dim) {
-      RAFT_CUDA_TRY(cudaMemset2DAsync(destination + preflight.dim,


The sorter still computes distances over the full row stride, so removing this makes nonzero padding affect which edges survive, no? Can we make it use the logical dimension separately from the stride? A test with dim 17 and stride 20 and nonzero padding would catch this.

dantegd · 2026-09-28T22:55:42Z

+
+  auto surviving_rows = raft::make_device_csr_matrix<uint32_t, int64_t, int64_t, int64_t>(
+    handle, 1, static_cast<std::size_t>(unfiltered_offsets.back()));
+  surviving_rows.initialize_sparsity(final_rows);


Could we count retained rows per input instead of materializing every surviving row ID? For 100M survivors, this needs over 1GB on the device plus 0.8 GB on the host just to compute offsets. concatenate_and_filter_datasets repeats the enumeration when both helpers are used.

dantegd · 2026-09-28T22:57:00Z

+		}, nil, nil
+	}
+
+	bitset := createBitset(allowList)


Not a Go expert, so not sure if I'm correct. But with 512 rows and an allowlist containing only row 0, this allocates one word, but MergedDatasetOffsets reads 16. We need zero-filled trailing words when the filter drops the final rows. Can we size this from the total input row count?

…a, with native BITSET-filtered merge support Java's merge() had no way to pass a bitset filter to the native cuvsCagraMerge call, and no public offsets helper -- the existing testFilteredMerge() test only worked by filtering rows in a Java for loop and calling the unfiltered merge() with the pre-filtered buffer. This closes that gap: - CagraIndex.merge() gains two new overloads taking a BitSet filter (a set bit keeps the row), threaded through to cuvsCagraMerge's native BITSET filter path instead of being emulated on the host. - New CagraIndex.mergedDatasetOffsets(indexes, filter), wrapping cuvsCagraMergedDatasetOffsets, since per-index offsets under a bitset filter can't be derived any other way. - New CagraIndex.concatenateDatasets()/concatenateAndFilterDatasets(), wrapping the new cuvsCagraConcatenateDatasets/ cuvsCagraConcatenateAndFilterDatasets C functions, as convenience helpers for building merge()'s mergedDataset argument. All four new CuVSProvider SPI methods are `default`, throwing UnsupportedOperationException (or delegating to the existing unfiltered overload for filter == null), so adding them doesn't break binary compatibility with out-of-tree providers. CagraIndexImpl.uploadBitsetFilter() uploads a Java BitSet (64-bit words) as the uint32-word device bitset cuVS expects via a direct byte-reinterpretation (a 64-bit little-endian word is exactly two adjacent uint32 words), reusing the same buildMemorySegment() zero-padding helper already used for prefilters elsewhere. Adds testConcatenateDatasetsMerge and testConcatenateAndFilterDatasetsMerge, covering both new helpers plus the native filtered-merge path end to end. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

HowardHuang1 · 2026-09-29T05:08:31Z

/ok to test 7df7f51

cjnolet · 2026-09-29T20:03:27Z

+ *                            it with `cuvsDatasetDestroy` when done.
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsCagraConcatenateAndFilterDatasets(cuvsResources_t res,


Let's simplify this and just have 1 function for concatenating which accepts a filter (aka consolidate this and the function above).

cjnolet · 2026-09-29T20:04:58Z

+ *                            `cuvsDatasetDestroy` when done.
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsCagraConcatenateDatasets(cuvsResources_t res,


I don't see why this is CAGRA-specific. This looks like a dataset-specific function. Please move to dataset.h.

…compat The PR renamed these functions to take caller-owned merged_dataset + offsets, breaking the C ABI. Per the ABI versioning convention, copy the new signatures to cuvsCagraMerge_v2 / cuvsCagraMergeWithParams_v2 and restore the original names with their pre-PR signatures, marked deprecated. The deprecated functions internally delegate to the new helpers: cuvsCagraConcatenateDatasets / cuvsCagraConcatenateAndFilterDatasets to build the merged buffer, cuvsCagraMergedDatasetOffsets for per-index offsets, then cuvsCagraMergeWithParams_v2 to merge the graph. Dataset ownership is transferred into the caller's empty handle via shallow-copy + wrapper delete. All new-signature call sites (C tests, Java CagraIndexImpl, Java headers_h Panama bindings) updated to use the _v2 names. The _v2 suffix will be dropped and the deprecated functions removed in 27.02. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

HowardHuang1 · 2026-10-01T06:57:00Z

/ok to test 3d8592b

…ergeWithParams_v2 These bindings already passed the new-style arguments (merged_dataset as IN, offsets array), but were still calling the deprecated old-named functions. The deprecated wrappers expect the old OUT-parameter contract so they either reject the call at runtime (Rust) or produce a type mismatch at compile time (Python cagra.cxx: int64_t* vs cuvsCagraIndex_t, Go: too many arguments). Fix: rename all call sites and declarations in Go, Python (.pxd + .pyx) and Rust (cuvs-sys/bindings.rs + cuvs/index.rs) to use the _v2 symbols. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

HowardHuang1 requested review from a team as code owners August 26, 2026 04:26

HowardHuang1 self-assigned this Aug 26, 2026

HowardHuang1 added breaking Introduces a breaking change Enhancement labels Aug 26, 2026

HowardHuang1 force-pushed the HH-migrate-merge-dataset-to-take-inputs-concat-buffer-and-offset-26_10 branch from 46bed89 to 87a0075 Compare August 26, 2026 04:33

HowardHuang1 added the improvement Improves an existing functionality label Aug 26, 2026

Merge remote-tracking branch 'upstream/main' into HH-migrate-merge-da…

50dd606

…taset-to-take-inputs-concat-buffer-and-offset-26_10

HowardHuang1 and others added 3 commits August 27, 2026 10:21

Merge remote-tracking branch 'upstream/main' into HH-migrate-merge-da…

3210112

…taset-to-take-inputs-concat-buffer-and-offset-26_10

Merge remote-tracking branch 'upstream/main' into HH-migrate-merge-da…

a83394f

…taset-to-take-inputs-concat-buffer-and-offset-26_10

HowardHuang1 changed the base branch from main to release/26.10 September 8, 2026 22:43

HowardHuang1 requested a review from a team as a code owner September 9, 2026 06:19

HowardHuang1 force-pushed the HH-migrate-merge-dataset-to-take-inputs-concat-buffer-and-offset-26_10 branch from 43620ae to a83394f Compare September 9, 2026 22:54

HowardHuang1 added 3 commits September 10, 2026 15:52

Merge remote-tracking branch 'upstream/release/26.10' into HH-migrate…

512e8b8

…-merge-dataset-to-take-inputs-concat-buffer-and-offset-26_10

Merge remote-tracking branch 'upstream/release/26.10' into HH-migrate…

8054382

…-merge-dataset-to-take-inputs-concat-buffer-and-offset-26_10

aamijar reviewed Sep 14, 2026

View reviewed changes

HowardHuang1 added 2 commits September 14, 2026 15:26

remove kvikio dependency by using public build() function rather than…

250b4da

… private detail::build_from_device_matrix()

Fix failing cuvs lucene tests

7a4845d

add fern changes from fern pre-commit hooks so that style-checker on …

a14a1da

…CI can pass

Merge remote-tracking branch 'upstream/release/26.10' into HH-migrate…

ab5fe77

…-merge-dataset-to-take-inputs-concat-buffer-and-offset-26_10

Regenerate Fern Lucene writer API docs after merge line-number shift

a8e86d7

cjnolet reviewed Sep 21, 2026

View reviewed changes

HowardHuang1 and others added 3 commits September 22, 2026 21:15

Merge remote-tracking branch 'upstream/release/26.10' into HH-migrate…

cf27fa0

…-merge-dataset-to-take-inputs-concat-buffer-and-offset-26_10

aamijar added this to Unstructured Data Processing Sep 24, 2026

aamijar moved this to Needs Review in Unstructured Data Processing Sep 24, 2026

dantegd requested changes Sep 28, 2026

View reviewed changes

cjnolet reviewed Sep 29, 2026

View reviewed changes

Conversation

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

Uh oh!

Uh oh!

copy-pr-bot Bot commented Aug 26, 2026

Uh oh!

HowardHuang1 commented Aug 26, 2026

Uh oh!

HowardHuang1 commented Aug 26, 2026

Uh oh!

HowardHuang1 commented Aug 27, 2026

Uh oh!

HowardHuang1 commented Sep 14, 2026

Uh oh!

aamijar left a comment

Choose a reason for hiding this comment

Uh oh!

HowardHuang1 commented Sep 14, 2026

Uh oh!

HowardHuang1 commented Sep 15, 2026

Uh oh!

HowardHuang1 commented Sep 18, 2026

Uh oh!

HowardHuang1 commented Sep 20, 2026

Uh oh!

HowardHuang1 commented Sep 21, 2026

Uh oh!

cjnolet Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

dantegd Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

dantegd Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

dantegd Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

dantegd Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

HowardHuang1 commented Sep 29, 2026

Uh oh!

cjnolet Sep 29, 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!

cjnolet Sep 29, 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!

HowardHuang1 commented Oct 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

HowardHuang1 commented Aug 26, 2026 •

edited

Loading

cjnolet Sep 29, 2026 •

edited

Loading

cjnolet Sep 29, 2026 •

edited

Loading