[Java] Encapsulate on-heap float arrays into `Dataset` #1024

ldematte · 2025-06-16T13:36:16Z

This PR is a follow-up from #902.
Still WIP (see self-comments on the changes) but I'd like some early feedback.

copy-pr-bot · 2025-06-16T13:36:19Z

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.

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ArrayDatasetImpl.java

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java

ldematte · 2025-06-19T08:07:03Z

@chatman @mythrocks wdyt?

ldematte · 2025-06-20T13:38:42Z

I added another commit 3f348c6 in which I reduce the scope for the Dataset/MemorySegment: I think there is no need to keep it around once the index is built, it is not referenced again (but keep me honest here and point out if I'm wrong!)

mythrocks · 2025-06-20T14:17:40Z

@chatman @mythrocks wdyt?

Looking forward to reading/reviewing this next week.

chatman · 2025-06-20T14:53:23Z

Thanks @ldematte. Looks good. I'll take a closer look at the scope of the memory segment. I remember extending to scope of the memory segment to closing it at index close due to some issues otherwise. I'll take another look today to make sure that's not happening.

…

On Fri, 20 Jun, 2025, 7:48 pm MithunR, ***@***.***> wrote: *mythrocks* left a comment (rapidsai/cuvs#1024) <#1024 (comment)> @chatman <https://github.com/chatman> @mythrocks <https://github.com/mythrocks> wdyt? Looking forward to reading/reviewing this next week. — Reply to this email directly, view it on GitHub <#1024 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABDCR5EJPEHUHFBOVUN3AY33EQJZXAVCNFSM6AAAAAB7NK52OWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDSOJRHAYDEMBXG4> . You are receiving this because you were mentioned.Message ID: ***@***.***>

ldematte · 2025-06-20T15:08:49Z

I'll take a closer look at the scope of the
memory segment. I remember extending to scope of the memory segment to
closing it at index close due to some issues otherwise. I'll take another
look today to make sure that's not happening.

Yes please; it might be that the C++ side of things holds a reference to the input memory (wrapped and managed by the MemorySegment), and therefore that that memory needs to be available for the lifetime of the Index.
If we even suspect this is the case, I'll revert that change - better hold onto some memory a bit longer than shorter than we need!

ldematte · 2025-06-20T15:10:19Z

BTW, no hurry -- I wasn't aware you guys were off this week. This can wait till next week, no prob.

ldematte · 2025-06-26T16:19:04Z

I'll take a closer look at the scope of the
memory segment. I remember extending to scope of the memory segment to
closing it at index close due to some issues otherwise.

Hey @chatman, I dug into the C and C++ implementations of cuvsCagraBuild/cuvs::neighbors::cagra::build and a pointer to the original dataset seems to be passed down and retained/referenced by Index. Seems because the DLTensor gets transformed to a mdspan, but I think it's just wrapped and no copy happens there, so we need to retain the dataset memorysegment "alive" for the lifetime of the index. I'll change this PR to do that.

chatman · 2025-06-26T16:39:32Z

Thanks a lot for looking deeper. This was precisely my concern, and thanks for looking into it. @cjnolet might know for sure, but it would be pretty safe to keep the scope tied up with the lifecycle of the index.

…

On Thu, 26 Jun, 2025, 9:49 pm Lorenzo Dematté, ***@***.***> wrote: *ldematte* left a comment (rapidsai/cuvs#1024) <#1024 (comment)> I'll take a closer look at the scope of the memory segment. I remember extending to scope of the memory segment to closing it at index close due to some issues otherwise. Hey @chatman <https://github.com/chatman>, I dug into the C and C++ implementations of cuvsCagraBuild/cuvs::neighbors::cagra::build and a pointer to the original dataset seems to be passed down and retained/referenced by Index. *Seems* because the DLTensor gets transformed to a mdspan, but I think it's just wrapped and no copy happens there, so we need to retain the dataset memorysegment "alive" for the lifetime of the index. I'll change this PR to do that. — Reply to this email directly, view it on GitHub <#1024 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABDCR5DADMYMRXZJVYNZUX33FQMQ5AVCNFSM6AAAAAB7NK52OWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTAMBZGA2DCNBWGA> . You are receiving this because you were mentioned.Message ID: ***@***.***>

…fferently

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java

ldematte · 2025-06-30T13:21:03Z

but it would be pretty safe to keep the scope tied up with the lifecycle of the index.

👍
I've done exactly that in the last commit

…array-dataset

ChrisHegarty

LGTM. This is a nice simplification.

mythrocks · 2025-07-03T21:07:54Z

java/README.md

+mvn clean integration-test -Dit.test=com.nvidia.cuvs.CagraBuildAndSearchIT
+```
+or, for a single test:
+```sh
+mvn clean integration-test -Dit.test=com.nvidia.cuvs.CagraBuildAndSearchIT#testMergeStrategies


java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java

mythrocks

LGTM! Will schedule to merge this once #1045 is merged.

…array-dataset # Conflicts: # java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java # java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java

ldematte · 2025-07-04T07:39:28Z

@mythrocks rebased on top of #1045 and ready to go! 👍

mythrocks · 2025-07-04T16:26:58Z

/ok to test 3b36197

mythrocks · 2025-07-04T21:08:50Z

/ok to test b0b3552

mythrocks · 2025-07-04T21:10:06Z

I just resolved the conflicts from #1081. This should go in after CI passes.

mythrocks · 2025-07-06T01:43:53Z

/merge

mythrocks · 2025-07-06T01:44:35Z

Thank you for this change, @ldematte. This has been merged.

This PR adds the ability to define a Dataset directly over a MemorySegment, "wrapping" it instead of allocating a new one. - Depends on #1033 and #1024 - ~~The new API has a `Object memorySegment` parameter, as we target Java 21 for the API (but 22 for the implementation); it works but it's definitely a hack and we need to sort this out~~ - As discussed, we want to keep targeting Java 21 for the API. This means the API will return a `MethodHandle`, and the Java 22 implementation will use it to return a factory method to build a Dataset from a MemorySegment. - This factory method can then be used as shown in the tests (see the `DatasetHelper` convenience class/method). - Benchmarks show a sizeable speedup -- it is still tiny related to the "big picture" (index build time), but there is an improvement and above all we avoid a whole new copy of the input data (halving the memory requirements). Fixes #698 Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) - Ben Frederickson (https://github.com/benfred) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: #1034

This PR adds the ability to define a Dataset directly over a MemorySegment, "wrapping" it instead of allocating a new one. - Depends on rapidsai#1033 and rapidsai#1024 - ~~The new API has a `Object memorySegment` parameter, as we target Java 21 for the API (but 22 for the implementation); it works but it's definitely a hack and we need to sort this out~~ - As discussed, we want to keep targeting Java 21 for the API. This means the API will return a `MethodHandle`, and the Java 22 implementation will use it to return a factory method to build a Dataset from a MemorySegment. - This factory method can then be used as shown in the tests (see the `DatasetHelper` convenience class/method). - Benchmarks show a sizeable speedup -- it is still tiny related to the "big picture" (index build time), but there is an improvement and above all we avoid a whole new copy of the input data (halving the memory requirements). Fixes rapidsai#698 Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) - Ben Frederickson (https://github.com/benfred) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: rapidsai#1034

This PR is a follow-up from rapidsai#902. Still WIP (see self-comments on the changes) but I'd like some early feedback. Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: rapidsai#1024

This PR adds the ability to define a Dataset directly over a MemorySegment, "wrapping" it instead of allocating a new one. - Depends on rapidsai#1033 and rapidsai#1024 - ~~The new API has a `Object memorySegment` parameter, as we target Java 21 for the API (but 22 for the implementation); it works but it's definitely a hack and we need to sort this out~~ - As discussed, we want to keep targeting Java 21 for the API. This means the API will return a `MethodHandle`, and the Java 22 implementation will use it to return a factory method to build a Dataset from a MemorySegment. - This factory method can then be used as shown in the tests (see the `DatasetHelper` convenience class/method). - Benchmarks show a sizeable speedup -- it is still tiny related to the "big picture" (index build time), but there is an improvement and above all we avoid a whole new copy of the input data (halving the memory requirements). Fixes rapidsai#698 Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) - Ben Frederickson (https://github.com/benfred) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: rapidsai#1034

This PR tidies up the lifecycle of `MemorySegment`s by using specific confined `Arena`s where possible, and specific index-bound `Arena`s where the lifetime needs to be bound to the lifetime of an index. It addresses all remaining usages of `CuVSResources#arena` after #1024 and #1045; as such, we can remove `CuVSResources#arena` completely, and force native memory usages to be tracked and dealt with specifically. This would towards the goal of #1037 Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) - Ben Frederickson (https://github.com/benfred) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: #1069

This PR tidies up the lifecycle of `MemorySegment`s by using specific confined `Arena`s where possible, and specific index-bound `Arena`s where the lifetime needs to be bound to the lifetime of an index. It addresses all remaining usages of `CuVSResources#arena` after rapidsai#1024 and rapidsai#1045; as such, we can remove `CuVSResources#arena` completely, and force native memory usages to be tracked and dealt with specifically. This would towards the goal of rapidsai#1037 Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) - Ben Frederickson (https://github.com/benfred) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: rapidsai#1069

This PR is a follow-up from rapidsai#902. Still WIP (see self-comments on the changes) but I'd like some early feedback. Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: rapidsai#1024

This PR adds the ability to define a Dataset directly over a MemorySegment, "wrapping" it instead of allocating a new one. - Depends on rapidsai#1033 and rapidsai#1024 - ~~The new API has a `Object memorySegment` parameter, as we target Java 21 for the API (but 22 for the implementation); it works but it's definitely a hack and we need to sort this out~~ - As discussed, we want to keep targeting Java 21 for the API. This means the API will return a `MethodHandle`, and the Java 22 implementation will use it to return a factory method to build a Dataset from a MemorySegment. - This factory method can then be used as shown in the tests (see the `DatasetHelper` convenience class/method). - Benchmarks show a sizeable speedup -- it is still tiny related to the "big picture" (index build time), but there is an improvement and above all we avoid a whole new copy of the input data (halving the memory requirements). Fixes rapidsai#698 Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) - Ben Frederickson (https://github.com/benfred) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: rapidsai#1034

ldematte commented Jun 16, 2025

View reviewed changes

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ArrayDatasetImpl.java Outdated Show resolved Hide resolved

ldematte commented Jun 16, 2025

View reviewed changes

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java Outdated Show resolved Hide resolved

ldematte added 2 commits June 16, 2025 17:31

Use Dataset to encapsulate on-heap arrays too

d0f53c0

Missing assert; renaming

114285a

ldematte changed the title ~~[WIP][Java] Encapsulate on-heap float arrays into Dataset~~ [Java] Encapsulate on-heap float arrays into Dataset Jun 19, 2025

ldematte mentioned this pull request Jun 19, 2025

[Java] Add Dataset based on MemorySegment #1034

Merged

ldematte mentioned this pull request Jun 20, 2025

[WIP][Java] Use MemorySegment for index serialization #1035

Closed

ldematte added 2 commits June 20, 2025 17:35

Different MemorySegment lifetime management (reduced scope)

3f348c6

Move addVector to Dataset.Builder

97c01ed

ldematte added 2 commits June 26, 2025 18:50

Extending Dataset lifetime to coincide with Index lifetime

052b62f

Reverting dataset change for BruteForceIndex -- it will be handled di…

47cc5bc

…fferently

ldematte commented Jun 26, 2025

View reviewed changes

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java Outdated Show resolved Hide resolved

ldematte commented Jun 26, 2025

View reviewed changes

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java Show resolved Hide resolved

ldematte commented Jun 26, 2025

View reviewed changes

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java Outdated Show resolved Hide resolved

ldematte changed the title ~~[Java] Encapsulate on-heap float arrays into Dataset~~ [Review][Java] Encapsulate on-heap float arrays into Dataset Jun 26, 2025

Merge branch 'branch-25.08' into java/float-array-dataset

c54faba

This was referenced Jun 27, 2025

[Improvement] [Java] Audit memory usage in the Java API #1037

Open

[Java] Tidy up MemorySegments lifecycle #1069

Merged

Merge remote-tracking branch 'upstream/branch-25.08' into java/float-…

30f4441

…array-dataset

ChrisHegarty approved these changes Jul 3, 2025

View reviewed changes

mythrocks reviewed Jul 3, 2025

View reviewed changes

java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java Show resolved Hide resolved

mythrocks approved these changes Jul 3, 2025

View reviewed changes

ldematte added 2 commits July 4, 2025 09:31

Merge remote-tracking branch 'upstream/branch-25.08' into java/float-…

381f74b

…array-dataset # Conflicts: # java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java # java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java

Protect against Dataset double-close

3b36197

mythrocks assigned ldematte Jul 4, 2025

mythrocks added improvement Improves an existing functionality non-breaking Introduces a non-breaking change Java labels Jul 4, 2025

mythrocks changed the title ~~[Review][Java] Encapsulate on-heap float arrays into Dataset~~ [Java] Encapsulate on-heap float arrays into Dataset Jul 4, 2025

Merge branch 'branch-25.08' into java/float-array-dataset

b0b3552

rapids-bot bot merged commit f02ee48 into rapidsai:branch-25.08 Jul 6, 2025
100 of 102 checks passed

[Java] Encapsulate on-heap float arrays into Dataset #1024

[Java] Encapsulate on-heap float arrays into Dataset #1024

Uh oh!

Conversation

ldematte commented Jun 16, 2025

Uh oh!

copy-pr-bot bot commented Jun 16, 2025

Uh oh!

Uh oh!

Uh oh!

ldematte commented Jun 19, 2025

Uh oh!

ldematte commented Jun 20, 2025

Uh oh!

mythrocks commented Jun 20, 2025

Uh oh!

chatman commented Jun 20, 2025 via email

Uh oh!

ldematte commented Jun 20, 2025

Uh oh!

ldematte commented Jun 20, 2025

Uh oh!

ldematte commented Jun 26, 2025

Uh oh!

chatman commented Jun 26, 2025 via email

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ldematte commented Jun 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ChrisHegarty left a comment

Choose a reason for hiding this comment

Uh oh!

mythrocks Jul 3, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mythrocks left a comment

Choose a reason for hiding this comment

Uh oh!

ldematte commented Jul 4, 2025

Uh oh!

mythrocks commented Jul 4, 2025

Uh oh!

mythrocks commented Jul 4, 2025

Uh oh!

mythrocks commented Jul 4, 2025

Uh oh!

mythrocks commented Jul 6, 2025

Uh oh!

Uh oh!

mythrocks commented Jul 6, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

[Java] Encapsulate on-heap float arrays into `Dataset` #1024

[Java] Encapsulate on-heap float arrays into `Dataset` #1024

ldematte commented Jun 30, 2025 •

edited

Loading