Forward-merge branch-25.08 into branch-25.10

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

@@ -77,12 +77,12 @@ private static MemorySegment createSampleDataSegment(Arena arena, float[][] arra
         return segment;
     }
 
-    private static Dataset fromMemorySegment(MemorySegment memorySegment, int size, int dimensions) {
+    private static CuVSMatrix fromMemorySegment(MemorySegment memorySegment, int size, int dimensions) {
         try {
-            return (Dataset)
+            return (CuVSMatrix)
                     CuVSProvider.provider()
-                            .newNativeDatasetBuilder()
-                            .invokeExact(memorySegment, size, dimensions);
+                            .newNativeMatrixBuilder()
+                            .invokeExact(memorySegment, size, dimensions, CuVSMatrix.DataType.FLOAT);
         } catch (Throwable e) {
             if (e instanceof Error err) {
                 throw err;
@@ -185,7 +185,7 @@ public void testIndexingFromMemorySegment(Blackhole blackhole) throws Throwable
     @OutputTimeUnit(TimeUnit.NANOSECONDS)
     @BenchmarkMode(Mode.AverageTime)
     public void testDatasetFromHeap(Blackhole blackhole) throws Throwable {
-        try (var dataset = Dataset.ofArray(arrayDataset)) {
+        try (var dataset = CuVSMatrix.ofArray(arrayDataset)) {
             blackhole.consume(dataset);
         }
     }

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

@@ -111,10 +111,10 @@ interface Builder {
     /**
      * Sets the dataset for building the {@link BruteForceIndex}.
      *
-     * @param dataset a {@link Dataset} object containing the vectors
+     * @param dataset a {@link CuVSMatrix} object containing the vectors
      * @return an instance of this Builder
      */
-    Builder withDataset(Dataset dataset);
+    Builder withDataset(CuVSMatrix dataset);
 
     /**
      * Builds and returns an instance of {@link BruteForceIndex}.

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

@@ -219,10 +219,10 @@ interface Builder {
     /**
      * Sets the dataset for building the {@link CagraIndex}.
      *
-     * @param dataset a {@link Dataset} object containing the vectors
+     * @param dataset a {@link CuVSMatrix} object containing the vectors
      * @return an instance of this Builder
      */
-    Builder withDataset(Dataset dataset);
+    Builder withDataset(CuVSMatrix dataset);
 
     /**
      * Registers an instance of configured {@link CagraIndexParams} with this

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

@@ -0,0 +1,21 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.nvidia.cuvs;
+
+/**
+ * A Dataset implementation backed by device (GPU) memory.
+ */
+public interface CuVSDeviceMatrix extends CuVSMatrix {}

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.nvidia.cuvs;
+
+/**
+ * A Dataset implementation backed by host (CPU) memory.
+ */
+public interface CuVSHostMatrix extends CuVSMatrix {
+  int get(int row, int col);
+}

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

@@ -0,0 +1,156 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.nvidia.cuvs;
+
+import com.nvidia.cuvs.spi.CuVSProvider;
+
+/**
+ * This represents a wrapper for a dataset to be used for index construction.
+ * The purpose is to allow a caller to place the vectors into native memory
+ * directly, instead of requiring the caller to load all the vectors into the heap
+ * (e.g. with a float[][]).
+ *
+ * @since 25.06
+ */
+public interface CuVSMatrix extends AutoCloseable {
+
+  enum DataType {
+    FLOAT,
+    INT,
+    BYTE
+  }
+
+  enum MemoryKind {
+    HOST,
+    DEVICE
+  }
+
+  /**
+   * Creates a dataset from an on-heap array of vectors.
+   * This method will allocate an additional MemorySegment to hold the graph data.
+   *
+   * @since 25.08
+   */
+  static CuVSMatrix ofArray(float[][] vectors) {
+    return CuVSProvider.provider().newMatrixFromArray(vectors);
+  }
+
+  /**
+   * Creates a dataset from an on-heap array of vectors.
+   * This method will allocate an additional MemorySegment to hold the graph data.
+   *
+   * @since 25.08
+   */
+  static CuVSMatrix ofArray(int[][] vectors) {
+    return CuVSProvider.provider().newMatrixFromArray(vectors);
+  }
+
+  /**
+   * Creates a dataset from an on-heap array of vectors.
+   * This method will allocate an additional MemorySegment to hold the graph data.
+   *
+   * @since 25.08
+   */
+  static CuVSMatrix ofArray(byte[][] vectors) {
+    return CuVSProvider.provider().newMatrixFromArray(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);
+
+    /**
+     * Add a single vector to the dataset.
+     *
+     * @param vector A byte array of as many elements as the dimensions
+     */
+    void addVector(byte[] vector);
+
+    /**
+     * Add a single vector to the dataset.
+     *
+     * @param vector A int array of as many elements as the dimensions
+     */
+    void addVector(int[] vector);
+
+    CuVSMatrix build();
+  }
+
+  /**
+   * Returns a builder to create a new instance of a dataset
+   *
+   * @param size     Number of vectors in the dataset
+   * @param columns  Size of each vector in the dataset
+   * @param dataType The data type of the dataset elements
+   * @return new instance of {@link CuVSMatrix}
+   */
+  static CuVSMatrix.Builder builder(int size, int columns, DataType dataType) {
+    return CuVSProvider.provider().newMatrixBuilder(size, columns, dataType);
+  }
+
+  /**
+   * Gets the size of the dataset
+   *
+   * @return Size of the dataset
+   */
+  long size();
+
+  /**
+   * Gets the number of columns in the Dataset (e.g. the dimensions of the vectors in this dataset,
+   * or the graph degree for the graph represented as a list of neighbours
+   *
+   * @return Dimensions of the vectors in the dataset
+   */
+  long columns();
+
+  /**
+   * Get a view (0-copy) of the row data, as a list of integers (32 bit)
+   *
+   * @param row the row for which to return the data
+   */
+  RowView getRow(long row);
+
+  /**
+   * Copies the content of this dataset to an on-heap Java matrix (array of arrays).
+   *
+   * @param array the destination array. Must be of length {@link CuVSMatrix#size()} or bigger,
+   *              and each element must be of length {@link CuVSMatrix#columns()} or bigger.
+   */
+  void toArray(int[][] array);
+
+  /**
+   * Copies the content of this dataset to an on-heap Java matrix (array of arrays).
+   *
+   * @param array the destination array. Must be of length {@link CuVSMatrix#size()} or bigger,
+   *              and each element must be of length {@link CuVSMatrix#columns()} or bigger.
+   */
+  void toArray(float[][] array);
+
+  /**
+   * Copies the content of this dataset to an on-heap Java matrix (array of arrays).
+   *
+   * @param array the destination array. Must be of length {@link CuVSMatrix#size()} or bigger,
+   *              and each element must be of length {@link CuVSMatrix#columns()} or bigger.
+   */
+  void toArray(byte[][] array);
+
+  @Override
+  void close();
+}