[Java] Add Dataset based on MemorySegment (rapidsai#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

Author: punAhuja

java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java

@@ -16,20 +16,25 @@
 
 package com.nvidia.cuvs;
 
+import com.nvidia.cuvs.spi.CuVSProvider;
 import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
 import org.openjdk.jmh.annotations.Param;
 import org.openjdk.jmh.annotations.Scope;
 import org.openjdk.jmh.annotations.Setup;
-import org.openjdk.jmh.annotations.TearDown;
 import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
 import org.openjdk.jmh.infra.Blackhole;
 
 import java.lang.foreign.*;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Random;
 import java.util.UUID;
+import java.util.concurrent.TimeUnit;
 
 @Fork(value = 1, warmups = 0)
 @State(Scope.Benchmark)
@@ -43,6 +48,10 @@ public class CagraIndexBenchmarks {
 
     private float[][] arrayDataset;
 
+    private Arena arena;
+
+    private MemorySegment memorySegmentDataset;
+
     private static final Random random = new Random();
 
     private static float[][] createSampleData(int size, int dimensions) {
@@ -55,9 +64,48 @@ private static float[][] createSampleData(int size, int dimensions) {
         return array;
     }
 
+    private static MemorySegment createSampleDataSegment(Arena arena, float[][] array, int size, int dimensions) {
+        final ValueLayout.OfFloat C_FLOAT = (ValueLayout.OfFloat) Linker.nativeLinker().canonicalLayouts().get("float");
+
+        MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long)size * dimensions, C_FLOAT);
+
+        var segment = arena.allocate(dataMemoryLayout);
+        for (int i = 0; i < size; ++i) {
+            var vector = array[i];
+            MemorySegment.copy(vector, 0, segment, C_FLOAT, (i * dimensions * C_FLOAT.byteSize()), dimensions);
+        }
+        return segment;
+    }
+
+    private static Dataset fromMemorySegment(MemorySegment memorySegment, int size, int dimensions) {
+        try {
+            return (Dataset)
+                    CuVSProvider.provider()
+                            .newNativeDatasetBuilder()
+                            .invokeExact(memorySegment, size, dimensions);
+        } catch (Throwable e) {
+            if (e instanceof Error err) {
+                throw err;
+            } else if (e instanceof RuntimeException re) {
+                throw re;
+            } else {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
     @Setup
     public void initialize() {
+        arena = Arena.ofShared();
         arrayDataset = createSampleData(size, dims);
+        memorySegmentDataset = createSampleDataSegment(arena, arrayDataset, size, dims);
+    }
+
+    @TearDown
+    public void cleanUp() {
+        if (arena != null) {
+            arena.close();
+        }
     }
 
     @Benchmark
@@ -111,4 +159,43 @@ public void testIndexingFromHeap(Blackhole blackhole) throws Throwable {
             blackhole.consume(index);
         }
     }
+
+    @Benchmark
+    public void testIndexingFromMemorySegment(Blackhole blackhole) throws Throwable {
+        try (CuVSResources resources = CuVSResources.create()) {
+            // Configure index parameters
+            CagraIndexParams indexParams = new CagraIndexParams.Builder()
+                .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT)
+                .withGraphDegree(1)
+                .withIntermediateGraphDegree(2)
+                .withNumWriterThreads(32)
+                .withMetric(CagraIndexParams.CuvsDistanceType.L2Expanded)
+                .build();
+
+            // Create the index with the dataset
+            CagraIndex index = CagraIndex.newBuilder(resources)
+                .withDataset(fromMemorySegment(memorySegmentDataset, size, dims))
+                .withIndexParams(indexParams)
+                .build();
+            blackhole.consume(index);
+        }
+    }
+
+    @Benchmark
+    @OutputTimeUnit(TimeUnit.NANOSECONDS)
+    @BenchmarkMode(Mode.AverageTime)
+    public void testDatasetFromHeap(Blackhole blackhole) throws Throwable {
+        try (var dataset = Dataset.ofArray(arrayDataset)) {
+            blackhole.consume(dataset);
+        }
+    }
+
+    @Benchmark
+    @OutputTimeUnit(TimeUnit.NANOSECONDS)
+    @BenchmarkMode(Mode.AverageTime)
+    public void testDatasetFromMemorySegment(Blackhole blackhole) throws Throwable {
+        try (var dataset = fromMemorySegment(memorySegmentDataset, size, dims)) {
+            blackhole.consume(dataset);
+        }
+    }
 }

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

@@ -28,7 +28,7 @@
 public interface Dataset extends AutoCloseable {
 
   /**
-   * Creates a dataset from a on-heap array of vectors
+   * Creates a dataset from an on-heap array of vectors
    *
    * @since 25.08
    */

java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java

@@ -22,6 +22,8 @@
 import com.nvidia.cuvs.Dataset;
 import com.nvidia.cuvs.HnswIndex;
 import com.nvidia.cuvs.ScalarQuantizer;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodType;
 import java.nio.file.Path;
 
 /**
@@ -53,6 +55,22 @@ default Path nativeLibraryPath() {
   /** Create a {@link Dataset.Builder} instance **/
   Dataset.Builder newDatasetBuilder(int size, int dimensions);
 
+  /**
+   * Returns the factory method used to build a Dataset from native memory.
+   * The factory method will have this signature: {@code Dataset createNativeDataset(memorySegment, size, dimensions)},
+   * where {@code memorySegment} is a {@code java.lang.foreign.MemorySegment} containing {@code int size} vectors of
+   * {@code int dimensions} length.
+   * <p>
+   * In order to expose this factory in a way that is compatible with Java 21, the factory method is returned as a
+   * {@link MethodHandle} with {@link MethodType} equal to
+   * {@code (Dataset.class, MemorySegment.class, int.class, int.class)}.
+   * The caller will need to invoke the factory via the {@link MethodHandle#invokeExact} method:
+   * {@code Dataset dataset = (Dataset)newNativeDatasetBuilder().invokeExact(memorySegment, size, dimensions)}
+   * </p>
+   * @return a MethodHandle which can be invoked to build a Dataset from a {@code MemorySegment}
+   */
+  MethodHandle newNativeDatasetBuilder();
+
   /** Create a {@link Dataset} backed by a on-heap array **/
   Dataset newArrayDataset(float[][] vectors);
 

java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java

@@ -21,6 +21,7 @@
 import com.nvidia.cuvs.Dataset;
 import com.nvidia.cuvs.HnswIndex;
 import com.nvidia.cuvs.ScalarQuantizer;
+import java.lang.invoke.MethodHandle;
 import java.nio.file.Path;
 
 /**
@@ -58,6 +59,11 @@ public Dataset.Builder newDatasetBuilder(int size, int dimensions) {
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public MethodHandle newNativeDatasetBuilder() {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public Dataset newArrayDataset(float[][] vectors) {
     throw new UnsupportedOperationException();

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

@@ -39,9 +39,8 @@ public class CuVSResourcesImpl implements CuVSResources {
   /**
    * Constructor that allocates the resources needed for cuVS
    *
-   * @throws Throwable exception thrown when native function is invoked
    */
-  public CuVSResourcesImpl(Path tempDirectory) throws Throwable {
+  public CuVSResourcesImpl(Path tempDirectory) {
     this.tempDirectory = tempDirectory;
     this.arena = Arena.ofShared();
     try (var localArena = Arena.ofConfined()) {

java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java

@@ -35,12 +35,32 @@
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemoryLayout;
 import java.lang.foreign.MemorySegment;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Objects;
 
 final class JDKProvider implements CuVSProvider {
 
+  private static final MethodHandle createNativeDataset$mh = createNativeDatasetBuilder();
+
+  static MethodHandle createNativeDatasetBuilder() {
+    try {
+      var lookup = MethodHandles.lookup();
+      var mt = MethodType.methodType(Dataset.class, MemorySegment.class, int.class, int.class);
+      return lookup.findStatic(JDKProvider.class, "createNativeDataset", mt);
+    } catch (NoSuchMethodException | IllegalAccessException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private static Dataset createNativeDataset(
+      MemorySegment memorySegment, int size, int dimensions) {
+    return new DatasetImpl(null, memorySegment, size, dimensions);
+  }
+
   @Override
   public CuVSResources newCuVSResources(Path tempDirectory) throws Throwable {
     Objects.requireNonNull(tempDirectory);
@@ -110,6 +130,11 @@ public Dataset build() {
     };
   }
 
+  @Override
+  public MethodHandle newNativeDatasetBuilder() {
+    return createNativeDataset$mh;
+  }
+
   @Override
   public Dataset newArrayDataset(float[][] vectors) {
     Objects.requireNonNull(vectors);