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

This PR is a follow-up from #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: #1024

Author: ldematte

java/README.md

@@ -25,10 +25,14 @@ do `./build.sh java` in the top level directory or just do `./build.sh` in this
 
 Run `./build.sh --run-java-tests` from this directory.
 
-To run a single test:
+To run a single test suite:
 ```sh
 cd cuvs-java/
-mvn verify -Dit.test=com.nvidia.cuvs.CagraBuildAndSearchIT
+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
 ```
 Be sure to set (manually, if needed) your `LD_LIBRARY_PATH` to include the directory with the appropriate (matching)
 version of `libcuvs.so`.

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

@@ -143,13 +143,6 @@ default void serializeToHNSW(OutputStream outputStream, Path tempFile) throws Th
    */
   void serializeToHNSW(OutputStream outputStream, Path tempFile, int bufferLength) throws Throwable;
 
-  /**
-   * Gets an instance of {@link CagraIndexParams}
-   *
-   * @return an instance of {@link CagraIndexParams}
-   */
-  CagraIndexParams getCagraIndexParameters();
-
   /**
    * Gets an instance of {@link CuVSResources}
    *

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

@@ -28,34 +28,47 @@
 public interface Dataset extends AutoCloseable {
 
   /**
-   * Add a single vector to the dataset.
+   * Creates a dataset from a on-heap array of vectors
    *
-   * @param vector A float array of as many elements as the dimensions
+   * @since 25.08
    */
-  public void addVector(float[] vector);
+  static Dataset ofArray(float[][] vectors) {
+    return CuVSProvider.provider().newArrayDataset(vectors);
+  }
+
+  interface Builder {
+    /**
+     * Add a single vector to the dataset.
+     *
+     * @param vector A float array of as many elements as the dimensions
+     */
+    void addVector(float[] vector);
+
+    Dataset build();
+  }
 
   /**
-   * Create a new instance of a dataset
+   * Returns a builder to create a new instance of a dataset
    *
    * @param size       Number of vectors in the dataset
    * @param dimensions Size of each vector in the dataset
    * @return new instance of {@link Dataset}
    */
-  static Dataset create(int size, int dimensions) {
-    return CuVSProvider.provider().newDataset(size, dimensions);
+  static Dataset.Builder builder(int size, int dimensions) {
+    return CuVSProvider.provider().newDatasetBuilder(size, dimensions);
   }
 
   /**
    * Gets the size of the dataset
    *
    * @return Size of the dataset
    */
-  public int size();
+  int size();
 
   /**
    * Gets the dimensions of the vectors in this dataset
    *
    * @return Dimensions of the vectors in the dataset
    */
-  public int dimensions();
+  int dimensions();
 }

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

@@ -49,8 +49,11 @@ default Path nativeLibraryPath() {
   /** Creates a new CuVSResources. */
   CuVSResources newCuVSResources(Path tempDirectory) throws Throwable;
 
-  /** Create a {@link Dataset} instance **/
-  Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException;
+  /** Create a {@link Dataset.Builder} instance **/
+  Dataset.Builder newDatasetBuilder(int size, int dimensions);
+
+  /** Create a {@link Dataset} backed by a on-heap array **/
+  Dataset newArrayDataset(float[][] vectors);
 
   /** Creates a new BruteForceIndex Builder. */
   BruteForceIndex.Builder newBruteForceIndexBuilder(CuVSResources cuVSResources)

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

@@ -53,7 +53,12 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable {
   }
 
   @Override
-  public Dataset newDataset(int size, int dimensions) throws UnsupportedOperationException {
+  public Dataset.Builder newDatasetBuilder(int size, int dimensions) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public Dataset newArrayDataset(float[][] vectors) {
     throw new UnsupportedOperationException();
   }
 }

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

@@ -47,7 +47,6 @@
 import com.nvidia.cuvs.CuVSResources;
 import com.nvidia.cuvs.Dataset;
 import com.nvidia.cuvs.SearchResults;
-import com.nvidia.cuvs.internal.common.Util;
 import com.nvidia.cuvs.internal.panama.cuvsFilter;
 import java.io.InputStream;
 import java.io.OutputStream;
@@ -70,11 +69,8 @@
  */
 public class BruteForceIndexImpl implements BruteForceIndex {
 
-  private final float[][] vectors;
-  private final Dataset dataset;
   private final CuVSResourcesImpl resources;
   private final IndexReference bruteForceIndexReference;
-  private final BruteForceIndexParams bruteForceIndexParams;
   private boolean destroyed;
 
   /**
@@ -87,15 +83,14 @@ public class BruteForceIndexImpl implements BruteForceIndex {
    *                              holding the index parameters
    */
   private BruteForceIndexImpl(
-      float[][] vectors,
-      Dataset dataset,
-      CuVSResourcesImpl resources,
-      BruteForceIndexParams bruteForceIndexParams) {
-    this.vectors = vectors;
-    this.dataset = dataset;
-    this.resources = resources;
-    this.bruteForceIndexParams = bruteForceIndexParams;
-    this.bruteForceIndexReference = build();
+      Dataset dataset, CuVSResourcesImpl resources, BruteForceIndexParams bruteForceIndexParams)
+      throws Exception {
+    Objects.requireNonNull(dataset);
+    try (dataset) {
+      this.resources = resources;
+      assert dataset instanceof DatasetImpl;
+      this.bruteForceIndexReference = build((DatasetImpl) dataset, bruteForceIndexParams);
+    }
   }
 
   /**
@@ -106,9 +101,6 @@ private BruteForceIndexImpl(
    */
   private BruteForceIndexImpl(InputStream inputStream, CuVSResourcesImpl resources)
       throws Throwable {
-    this.bruteForceIndexParams = null;
-    this.vectors = null;
-    this.dataset = null;
     this.resources = resources;
     this.bruteForceIndexReference = deserialize(inputStream);
   }
@@ -124,7 +116,7 @@ private void checkNotDestroyed() {
    * BRUTEFORCE index
    */
   @Override
-  public void destroyIndex() throws Throwable {
+  public void destroyIndex() {
     checkNotDestroyed();
     try {
       int returnValue = cuvsBruteForceIndexDestroy(bruteForceIndexReference.indexPtr);
@@ -141,7 +133,6 @@ public void destroyIndex() throws Throwable {
     } finally {
       destroyed = true;
     }
-    if (dataset != null) dataset.close();
   }
 
   /**
@@ -151,16 +142,13 @@ public void destroyIndex() throws Throwable {
    * @return an instance of {@link IndexReference} that holds the pointer to the
    *         index
    */
-  private IndexReference build() {
+  private IndexReference build(DatasetImpl dataset, BruteForceIndexParams bruteForceIndexParams) {
     try (var localArena = Arena.ofConfined()) {
-      long rows = dataset != null ? dataset.size() : vectors.length;
-      long cols = dataset != null ? dataset.dimensions() : (rows > 0 ? vectors[0].length : 0);
+      long rows = dataset.size();
+      long cols = dataset.dimensions();
 
       Arena arena = resources.getArena();
-      MemorySegment datasetMemSegment =
-          dataset != null
-              ? ((DatasetImpl) dataset).seg
-              : Util.buildMemorySegment(resources.getArena(), vectors);
+      MemorySegment datasetMemSegment = dataset.asMemorySegment();
 
       long cuvsResources = resources.getHandle();
 
@@ -248,7 +236,7 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable {
       long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension;
       long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topk;
       long distanceBytes = C_FLOAT_BYTE_SIZE * numQueries * topk;
-      long prefilterBytes = 0;
+      long prefilterBytes = 0; // size assigned later
 
       returnValue = cuvsRMMAlloc(cuvsResources, queriesD, queriesBytes);
       checkCuVSError(returnValue, "cuvsRMMAlloc");
@@ -430,7 +418,6 @@ public static BruteForceIndex.Builder newBuilder(CuVSResources cuvsResources) {
    */
   public static class Builder implements BruteForceIndex.Builder {
 
-    private float[][] vectors;
     private Dataset dataset;
     private final CuVSResourcesImpl cuvsResources;
     private BruteForceIndexParams bruteForceIndexParams;
@@ -479,7 +466,7 @@ public Builder from(InputStream inputStream) {
      */
     @Override
     public Builder withDataset(float[][] vectors) {
-      this.vectors = vectors;
+      this.dataset = Dataset.ofArray(vectors);
       return this;
     }
 
@@ -505,7 +492,7 @@ public BruteForceIndexImpl build() throws Throwable {
       if (inputStream != null) {
         return new BruteForceIndexImpl(inputStream, cuvsResources);
       } else {
-        return new BruteForceIndexImpl(vectors, dataset, cuvsResources, bruteForceIndexParams);
+        return new BruteForceIndexImpl(dataset, cuvsResources, bruteForceIndexParams);
       }
     }
   }