diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java index b8a77979b5..18d8495424 100644 --- a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java +++ b/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); } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceIndex.java index b21109dc9e..2b50a0b741 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceIndex.java +++ b/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}. diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index 8e6e2fb27d..84487857c8 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/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 diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java new file mode 100644 index 0000000000..23d29fd479 --- /dev/null +++ b/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 {} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java new file mode 100644 index 0000000000..237dfa39f7 --- /dev/null +++ b/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); +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java new file mode 100644 index 0000000000..0b892fbc5f --- /dev/null +++ b/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(); +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java deleted file mode 100644 index e13bed45d0..0000000000 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java +++ /dev/null @@ -1,74 +0,0 @@ -/* - * 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 Dataset extends AutoCloseable { - - /** - * Creates a dataset from an on-heap array of vectors - * - * @since 25.08 - */ - 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(); - } - - /** - * 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.Builder builder(int size, int dimensions) { - return CuVSProvider.provider().newDatasetBuilder(size, dimensions); - } - - /** - * Gets the size of the dataset - * - * @return Size of the dataset - */ - int size(); - - /** - * Gets the dimensions of the vectors in this dataset - * - * @return Dimensions of the vectors in the dataset - */ - int dimensions(); -} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/RowView.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/RowView.java new file mode 100644 index 0000000000..cbace82ae3 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/RowView.java @@ -0,0 +1,74 @@ +/* + * 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; + +/** + * Represent a contiguous list of integers (32-bit) backed by off-heap memory. + * + * @since 25.08 + */ +public interface RowView { + + long size(); + + /** + * Returns the integer element at the given position. Asserts that the + * data type of the dataset on top of which this view is instantiates is + * {@link CuVSMatrix.DataType#INT} + * + * @param index the element index + */ + int getAsInt(long index); + + /** + * Returns the integer element at the given position. Asserts that the + * data type of the dataset on top of which this view is instantiates is + * {@link CuVSMatrix.DataType#FLOAT} + * + * @param index the element index + */ + float getAsFloat(long index); + + /** + * Returns the integer element at the given position. Asserts that the + * data type of the dataset on top of which this view is instantiates is + * {@link CuVSMatrix.DataType#BYTE} + * + * @param index the element index + */ + byte getAsByte(long index); + + /** + * Copies the content of this row to an on-heap Java array. + * + * @param array the destination array. Must be of length {@link RowView#size()} or bigger. + */ + void toArray(int[] array); + + /** + * Copies the content of this row to an on-heap Java array. + * + * @param array the destination array. Must be of length {@link RowView#size()} or bigger. + */ + void toArray(float[] array); + + /** + * Copies the content of this row to an on-heap Java array. + * + * @param array the destination array. Must be of length {@link RowView#size()} or bigger. + */ + void toArray(byte[] array); +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java index a6bd7e1f6e..2726382f23 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java @@ -104,10 +104,10 @@ interface Builder { /** * Sets the dataset for building the TieredIndex. * - * @param dataset A {@link Dataset} instance containing the vectors + * @param dataset A {@link CuVSMatrix} instance containing the vectors * @return This Builder instance for method chaining */ - Builder withDataset(Dataset dataset); + Builder withDataset(CuVSMatrix dataset); /** * Registers TieredIndex parameters with this Builder. @@ -165,11 +165,11 @@ interface ExtendBuilder { /** * Sets the dataset to add to the existing index. * - * @param dataset A {@link Dataset} instance containing the new vectors to + * @param dataset A {@link CuVSMatrix} instance containing the new vectors to * add * @return This ExtendBuilder instance for method chaining */ - ExtendBuilder withDataset(Dataset dataset); + ExtendBuilder withDataset(CuVSMatrix dataset); /** * Executes the extend operation, adding the specified data to the index. diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index 4dd7fa8eb8..e528beff1a 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -18,8 +18,8 @@ import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraMergeParams; +import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; import com.nvidia.cuvs.TieredIndex; import java.lang.invoke.MethodHandle; @@ -52,27 +52,34 @@ default Path nativeLibraryPath() { /** Creates a new CuVSResources. */ CuVSResources newCuVSResources(Path tempDirectory) throws Throwable; - /** Create a {@link Dataset.Builder} instance **/ - Dataset.Builder newDatasetBuilder(int size, int dimensions); + /** Create a {@link CuVSMatrix.Builder} instance **/ + CuVSMatrix.Builder newMatrixBuilder(int size, int dimensions, CuVSMatrix.DataType dataType); /** * 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)}, + * The factory method will have this signature: + * {@code Dataset createNativeDataset(memorySegment, size, dimensions, dataType)}, * where {@code memorySegment} is a {@code java.lang.foreign.MemorySegment} containing {@code int size} vectors of - * {@code int dimensions} length. + * {@code int dimensions} length of type {@link CuVSMatrix.DataType}. *

* 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)}. + * {@code (Dataset.class, MemorySegment.class, int.class, int.class, Dataset.DataType.class)}. * The caller will need to invoke the factory via the {@link MethodHandle#invokeExact} method: - * {@code Dataset dataset = (Dataset)newNativeDatasetBuilder().invokeExact(memorySegment, size, dimensions)} + * {@code Dataset dataset = (Dataset)newNativeDatasetBuilder().invokeExact(memorySegment, size, dimensions, dataType)} *

- * @return a MethodHandle which can be invoked to build a Dataset from a {@code MemorySegment} + * @return a MethodHandle which can be invoked to build a Dataset from an external {@code MemorySegment} */ - MethodHandle newNativeDatasetBuilder(); + MethodHandle newNativeMatrixBuilder(); - /** Create a {@link Dataset} backed by a on-heap array **/ - Dataset newArrayDataset(float[][] vectors); + /** Create a {@link CuVSMatrix} from an on-heap array **/ + CuVSMatrix newMatrixFromArray(float[][] vectors); + + /** Create a {@link CuVSMatrix} from an on-heap array **/ + CuVSMatrix newMatrixFromArray(int[][] vectors); + + /** Create a {@link CuVSMatrix} from an on-heap array **/ + CuVSMatrix newMatrixFromArray(byte[][] vectors); /** Creates a new BruteForceIndex Builder. */ BruteForceIndex.Builder newBruteForceIndexBuilder(CuVSResources cuVSResources) diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 7e5d1a7e64..8f65bf7068 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -17,8 +17,8 @@ import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.HnswIndex; import com.nvidia.cuvs.TieredIndex; import java.lang.invoke.MethodHandle; @@ -60,17 +60,28 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable { } @Override - public Dataset.Builder newDatasetBuilder(int size, int dimensions) { + public CuVSMatrix.Builder newMatrixBuilder( + int size, int dimensions, CuVSMatrix.DataType dataType) { throw new UnsupportedOperationException(); } @Override - public MethodHandle newNativeDatasetBuilder() { + public MethodHandle newNativeMatrixBuilder() { throw new UnsupportedOperationException(); } @Override - public Dataset newArrayDataset(float[][] vectors) { + public CuVSMatrix newMatrixFromArray(float[][] vectors) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSMatrix newMatrixFromArray(int[][] vectors) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSMatrix newMatrixFromArray(byte[][] vectors) { throw new UnsupportedOperationException(); } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index 163205b830..f32e311f77 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -42,8 +42,8 @@ import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.BruteForceIndexParams; import com.nvidia.cuvs.BruteForceQuery; +import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.SearchResults; import com.nvidia.cuvs.internal.panama.cuvsFilter; import java.io.InputStream; @@ -81,13 +81,13 @@ public class BruteForceIndexImpl implements BruteForceIndex { * holding the index parameters */ private BruteForceIndexImpl( - Dataset dataset, CuVSResources resources, BruteForceIndexParams bruteForceIndexParams) + CuVSMatrix dataset, CuVSResources resources, BruteForceIndexParams bruteForceIndexParams) throws Exception { Objects.requireNonNull(dataset); try (dataset) { this.resources = resources; - assert dataset instanceof DatasetImpl; - this.bruteForceIndexReference = build((DatasetImpl) dataset, bruteForceIndexParams); + assert dataset instanceof CuVSMatrixBaseImpl; + this.bruteForceIndexReference = build((CuVSMatrixBaseImpl) dataset, bruteForceIndexParams); } } @@ -144,11 +144,12 @@ public void destroyIndex() { * @return an instance of {@link IndexReference} that holds the pointer to the * index */ - private IndexReference build(DatasetImpl dataset, BruteForceIndexParams bruteForceIndexParams) { + private IndexReference build( + CuVSMatrixBaseImpl dataset, BruteForceIndexParams bruteForceIndexParams) { long rows = dataset.size(); - long cols = dataset.dimensions(); + long cols = dataset.columns(); - MemorySegment datasetMemSegment = dataset.asMemorySegment(); + MemorySegment datasetMemSegment = dataset.memorySegment(); omp_set_num_threads(bruteForceIndexParams.getNumWriterThreads()); @@ -164,7 +165,7 @@ private IndexReference build(DatasetImpl dataset, BruteForceIndexParams bruteFor long[] datasetShape = {rows, cols}; var tensorDataArena = Arena.ofShared(); MemorySegment datasetTensor = - prepareTensor(tensorDataArena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 2, 1); + prepareTensor(tensorDataArena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 1); var returnValue = cuvsStreamSync(cuvsResources); checkCuVSError(returnValue, "cuvsStreamSync"); @@ -236,13 +237,13 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { long[] queriesShape = {numQueries, vectorDimension}; MemorySegment queriesTensor = - prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 2, 1); + prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 1); long[] neighborsShape = {numQueries, topk}; MemorySegment neighborsTensor = - prepareTensor(localArena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1); + prepareTensor(localArena, neighborsDP, neighborsShape, 0, 64, 2, 1); long[] distancesShape = {numQueries, topk}; MemorySegment distancesTensor = - prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 2, 1); + prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 1); MemorySegment prefilter = cuvsFilter.allocate(localArena); MemorySegment prefilterTensor; @@ -259,7 +260,7 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { cudaMemcpy(prefilterDP, prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); - prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 1, 2, 1); + prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 2, 1); cuvsFilter.type(prefilter, 2); cuvsFilter.addr(prefilter, prefilterTensor.address()); @@ -393,7 +394,7 @@ public static BruteForceIndex.Builder newBuilder(CuVSResources cuvsResources) { */ public static class Builder implements BruteForceIndex.Builder { - private Dataset dataset; + private CuVSMatrix dataset; private final CuVSResources cuvsResources; private BruteForceIndexParams bruteForceIndexParams; private InputStream inputStream; @@ -441,18 +442,18 @@ public Builder from(InputStream inputStream) { */ @Override public Builder withDataset(float[][] vectors) { - this.dataset = Dataset.ofArray(vectors); + this.dataset = CuVSMatrix.ofArray(vectors); return this; } /** * 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 */ @Override - public Builder withDataset(Dataset dataset) { + public Builder withDataset(CuVSMatrix dataset) { this.dataset = dataset; return this; } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 028296d270..4fc373e8ca 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -39,8 +39,8 @@ import com.nvidia.cuvs.CagraSearchParams; import com.nvidia.cuvs.CuVSIvfPqIndexParams; import com.nvidia.cuvs.CuVSIvfPqSearchParams; +import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.Dataset; import com.nvidia.cuvs.SearchResults; import com.nvidia.cuvs.internal.common.CloseableHandle; import com.nvidia.cuvs.internal.common.CompositeCloseableHandle; @@ -90,11 +90,11 @@ public class CagraIndexImpl implements CagraIndex { * @param resources an instance of {@link CuVSResources} */ private CagraIndexImpl( - CagraIndexParams indexParameters, Dataset dataset, CuVSResources resources) { + CagraIndexParams indexParameters, CuVSMatrix dataset, CuVSResources resources) { Objects.requireNonNull(dataset); this.resources = resources; - assert dataset instanceof DatasetImpl; - this.cagraIndexReference = build(indexParameters, (DatasetImpl) dataset); + assert dataset instanceof CuVSMatrixBaseImpl; + this.cagraIndexReference = build(indexParameters, (CuVSMatrixBaseImpl) dataset); } /** @@ -151,9 +151,9 @@ public void destroyIndex() throws Throwable { * @return an instance of {@link IndexReference} that holds the pointer to the * index */ - private IndexReference build(CagraIndexParams indexParameters, DatasetImpl dataset) { + private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImpl dataset) { long rows = dataset.size(); - long cols = dataset.dimensions(); + long cols = dataset.columns(); try (var indexParams = segmentFromIndexParams(indexParameters); var localArena = Arena.ofConfined()) { @@ -162,11 +162,10 @@ private IndexReference build(CagraIndexParams indexParameters, DatasetImpl datas int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 1; omp_set_num_threads(numWriterThreads); - MemorySegment dataSeg = dataset.asMemorySegment(); + MemorySegment dataSeg = dataset.memorySegment(); long[] datasetShape = {rows, cols}; - MemorySegment datasetTensor = - prepareTensor(localArena, dataSeg, datasetShape, 2, 32, 2, 2, 1); + MemorySegment datasetTensor = prepareTensor(localArena, dataSeg, datasetShape, 2, 32, 2, 1); var index = createCagraIndex(); @@ -253,13 +252,13 @@ public SearchResults search(CagraQuery query) throws Throwable { long[] queriesShape = {numQueries, vectorDimension}; MemorySegment queriesTensor = - prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 2, 1); + prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 1); long[] neighborsShape = {numQueries, topK}; MemorySegment neighborsTensor = - prepareTensor(localArena, neighborsDP, neighborsShape, 1, 32, 2, 2, 1); + prepareTensor(localArena, neighborsDP, neighborsShape, 1, 32, 2, 1); long[] distancesShape = {numQueries, topK}; MemorySegment distancesTensor = - prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 2, 1); + prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 1); var returnValue = cuvsStreamSync(cuvsRes); checkCuVSError(returnValue, "cuvsStreamSync"); @@ -294,7 +293,7 @@ public SearchResults search(CagraQuery query) throws Throwable { cudaMemcpy(prefilterDP, prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); - prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 1, 2, 1); + prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 2, 1); cuvsFilter.type(prefilter, 1); cuvsFilter.addr(prefilter, prefilterTensor.address()); @@ -680,7 +679,7 @@ private static CloseableHandle createMergeParamsSegment(CagraMergeParams mergePa */ public static class Builder implements CagraIndex.Builder { - private Dataset dataset; + private CuVSMatrix dataset; private CagraIndexParams cagraIndexParams; private final CuVSResources cuvsResources; private InputStream inputStream; @@ -697,12 +696,12 @@ public Builder from(InputStream inputStream) { @Override public Builder withDataset(float[][] vectors) { - this.dataset = Dataset.ofArray(vectors); + this.dataset = CuVSMatrix.ofArray(vectors); return this; } @Override - public Builder withDataset(Dataset dataset) { + public Builder withDataset(CuVSMatrix dataset) { this.dataset = dataset; return this; } @@ -729,7 +728,7 @@ public CagraIndexImpl build() throws Throwable { public static class IndexReference { private final MemorySegment memorySegment; - private final Dataset dataset; + private final CuVSMatrix dataset; /** * Constructs CagraIndexReference with an instance of MemorySegment passed as a @@ -742,7 +741,7 @@ public static class IndexReference { * to it so we can close it when the index is closed. * Can be null (e.g. from deserialization or merging) */ - private IndexReference(MemorySegment indexMemorySegment, Dataset dataset) { + private IndexReference(MemorySegment indexMemorySegment, CuVSMatrix dataset) { this.memorySegment = indexMemorySegment; this.dataset = dataset; } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixArenaImpl.java similarity index 54% rename from java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java rename to java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixArenaImpl.java index ee065e632b..1dc2350638 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/DatasetImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixArenaImpl.java @@ -15,22 +15,37 @@ */ package com.nvidia.cuvs.internal; -import com.nvidia.cuvs.Dataset; import java.lang.foreign.Arena; -import java.lang.foreign.MemorySegment; +import java.lang.foreign.SequenceLayout; +import java.lang.foreign.ValueLayout; import java.util.concurrent.atomic.AtomicReference; -public class DatasetImpl implements Dataset { +/** + * A Dataset implementation backed by host (CPU) memory. + * Memory is allocated and managed by the implementation, via a Java shared {@link Arena} + */ +public class CuVSHostMatrixArenaImpl extends CuVSHostMatrixImpl { private final AtomicReference arenaReference; - private final MemorySegment seg; - private final int size; - private final int dimensions; - public DatasetImpl(Arena arena, MemorySegment memorySegment, int size, int dimensions) { + public CuVSHostMatrixArenaImpl(long size, long columns, DataType dataType) { + this( + size, + columns, + dataType, + valueLayoutFromType(dataType), + sequenceLayoutFromType(size, columns, dataType), + Arena.ofShared()); + } + + private CuVSHostMatrixArenaImpl( + long size, + long columns, + DataType dataType, + ValueLayout valueLayout, + SequenceLayout layout, + Arena arena) { + super(arena.allocate(layout), size, columns, dataType, valueLayout, layout); this.arenaReference = new AtomicReference<>(arena); - this.seg = memorySegment; - this.size = size; - this.dimensions = dimensions; } @Override @@ -40,18 +55,4 @@ public void close() { arena.close(); } } - - @Override - public int size() { - return size; - } - - @Override - public int dimensions() { - return dimensions; - } - - public MemorySegment asMemorySegment() { - return seg; - } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java new file mode 100644 index 0000000000..a7096e207c --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java @@ -0,0 +1,198 @@ +/* + * 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.internal; + +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_CHAR; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; + +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.RowView; +import java.lang.foreign.MemoryLayout; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.SequenceLayout; +import java.lang.foreign.ValueLayout; +import java.lang.invoke.VarHandle; + +/** + * A Dataset implementation backed by host (CPU) memory. + */ +public class CuVSHostMatrixImpl extends CuVSMatrixBaseImpl implements CuVSHostMatrix { + private final ValueLayout valueLayout; + protected final VarHandle accessor$vh; + + public CuVSHostMatrixImpl( + MemorySegment memorySegment, long size, long columns, DataType dataType) { + this( + memorySegment, + size, + columns, + dataType, + valueLayoutFromType(dataType), + MemoryLayout.sequenceLayout(size * columns, valueLayoutFromType(dataType)) + .withByteAlignment(32)); + } + + protected CuVSHostMatrixImpl( + MemorySegment memorySegment, + long size, + long columns, + DataType dataType, + ValueLayout valueLayout, + MemoryLayout sequenceLayout) { + super(memorySegment, dataType, size, columns); + this.accessor$vh = sequenceLayout.varHandle(MemoryLayout.PathElement.sequenceElement()); + this.valueLayout = valueLayout; + } + + protected static ValueLayout valueLayoutFromType(DataType dataType) { + return switch (dataType) { + case FLOAT -> C_FLOAT; + case INT -> C_INT; + case BYTE -> C_CHAR; + }; + } + + protected static SequenceLayout sequenceLayoutFromType( + long size, long columns, DataType dataType) { + return MemoryLayout.sequenceLayout(size * columns, valueLayoutFromType(dataType)) + .withByteAlignment(32); + } + + @Override + public RowView getRow(long nodeIndex) { + var valueByteSize = valueLayout.byteSize(); + return new SliceRowView( + memorySegment.asSlice(nodeIndex * columns * valueByteSize, columns * valueByteSize), + columns, + valueLayout, + dataType, + valueByteSize); + } + + @Override + public void toArray(int[][] array) { + assert dataType == DataType.INT; + assert (array.length >= size) : "Input array is not large enough"; + assert (array.length == 0 || array[0].length >= columns) : "Input array is not large enough"; + var valueByteSize = valueLayout.byteSize(); + for (int r = 0; r < size; ++r) { + MemorySegment.copy( + memorySegment, valueLayout, r * columns * valueByteSize, array[r], 0, (int) columns); + } + } + + @Override + public void toArray(float[][] array) { + assert dataType == DataType.FLOAT; + assert (array.length >= size) : "Input array is not large enough"; + assert (array.length == 0 || array[0].length >= columns) : "Input array is not large enough"; + var valueByteSize = valueLayout.byteSize(); + for (int r = 0; r < size; ++r) { + MemorySegment.copy( + memorySegment, valueLayout, r * columns * valueByteSize, array[r], 0, (int) columns); + } + } + + @Override + public void toArray(byte[][] array) { + assert dataType == DataType.BYTE; + assert (array.length >= size) : "Input array is not large enough"; + assert (array.length == 0 || array[0].length >= columns) : "Input array is not large enough"; + var valueByteSize = valueLayout.byteSize(); + for (int r = 0; r < size; ++r) { + MemorySegment.copy( + memorySegment, valueLayout, r * columns * valueByteSize, array[r], 0, (int) columns); + } + } + + @Override + public void close() {} + + @Override + public int get(int row, int col) { + return (int) accessor$vh.get(memorySegment, 0L, (long) row * columns + col); + } + + public ValueLayout valueLayout() { + return valueLayout; + } + + private static class SliceRowView implements RowView { + private final MemorySegment memorySegment; + private final long size; + private final ValueLayout valueLayout; + private final DataType dataType; + private final long valueByteSize; + + SliceRowView( + MemorySegment slice, + long size, + ValueLayout valueLayout, + DataType dataType, + long valueByteSize) { + this.memorySegment = slice; + this.size = size; + this.valueLayout = valueLayout; + this.dataType = dataType; + this.valueByteSize = valueByteSize; + } + + @Override + public long size() { + return size; + } + + @Override + public float getAsFloat(long index) { + assert dataType == DataType.FLOAT; + return memorySegment.get((ValueLayout.OfFloat) valueLayout, index * valueByteSize); + } + + @Override + public byte getAsByte(long index) { + assert dataType == DataType.BYTE; + return memorySegment.get((ValueLayout.OfByte) valueLayout, index * valueByteSize); + } + + @Override + public int getAsInt(long index) { + assert dataType == DataType.INT; + return memorySegment.get((ValueLayout.OfInt) valueLayout, index * valueByteSize); + } + + @Override + public void toArray(int[] array) { + assert (array.length >= size) : "Input array is not large enough"; + assert dataType == DataType.INT; + MemorySegment.copy(memorySegment, valueLayout, 0, array, 0, (int) size); + } + + @Override + public void toArray(float[] array) { + assert (array.length >= size) : "Input array is not large enough"; + assert dataType == DataType.FLOAT; + MemorySegment.copy(memorySegment, valueLayout, 0, array, 0, (int) size); + } + + @Override + public void toArray(byte[] array) { + assert (array.length >= size) : "Input array is not large enough"; + assert dataType == DataType.BYTE; + MemorySegment.copy(memorySegment, valueLayout, 0, array, 0, (int) size); + } + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java new file mode 100644 index 0000000000..fc9daf9387 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java @@ -0,0 +1,48 @@ +/* + * 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.internal; + +import com.nvidia.cuvs.CuVSMatrix; +import java.lang.foreign.MemorySegment; + +public abstract class CuVSMatrixBaseImpl implements CuVSMatrix { + protected final MemorySegment memorySegment; + protected final DataType dataType; + protected final long size; + protected final long columns; + + protected CuVSMatrixBaseImpl( + MemorySegment memorySegment, DataType dataType, long size, long columns) { + this.memorySegment = memorySegment; + this.dataType = dataType; + this.size = size; + this.columns = columns; + } + + @Override + public long size() { + return size; + } + + @Override + public long columns() { + return columns; + } + + public MemorySegment memorySegment() { + return memorySegment; + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java index a4231450e2..1c97422d15 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java @@ -103,14 +103,13 @@ public SearchResults search(HnswQuery query) throws Throwable { MemorySegment querySeg = buildMemorySegment(localArena, queryVectors); long[] queriesShape = {numQueries, vectorDimension}; - MemorySegment queriesTensor = - prepareTensor(localArena, querySeg, queriesShape, 2, 32, 2, 1, 1); + MemorySegment queriesTensor = prepareTensor(localArena, querySeg, queriesShape, 2, 32, 1, 1); long[] neighborsShape = {numQueries, topK}; MemorySegment neighborsTensor = - prepareTensor(localArena, neighborsMemorySegment, neighborsShape, 1, 64, 2, 1, 1); + prepareTensor(localArena, neighborsMemorySegment, neighborsShape, 1, 64, 1, 1); long[] distancesShape = {numQueries, topK}; MemorySegment distancesTensor = - prepareTensor(localArena, distancesMemorySegment, distancesShape, 2, 32, 2, 1, 1); + prepareTensor(localArena, distancesMemorySegment, distancesShape, 2, 32, 1, 1); try (var resourcesAccessor = resources.access()) { var cuvsRes = resourcesAccessor.handle(); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java index f5dcb5eff0..9dfbf5a843 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java @@ -30,14 +30,8 @@ import static com.nvidia.cuvs.internal.common.Util.prepareTensor; import static com.nvidia.cuvs.internal.panama.headers_h.*; -import com.nvidia.cuvs.CagraIndexParams; -import com.nvidia.cuvs.CagraSearchParams; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.Dataset; -import com.nvidia.cuvs.SearchResults; -import com.nvidia.cuvs.TieredIndex; -import com.nvidia.cuvs.TieredIndexParams; -import com.nvidia.cuvs.TieredIndexQuery; +import com.nvidia.cuvs.*; +import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.internal.common.Util; import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams; import com.nvidia.cuvs.internal.panama.cuvsCagraSearchParams; @@ -62,7 +56,7 @@ * @since 25.02 */ public class TieredIndexImpl implements TieredIndex { - private final Dataset dataset; + private final CuVSMatrix dataset; private final CuVSResources resources; private final TieredIndexParams tieredIndexParameters; private final IndexReference tieredIndexReference; @@ -72,7 +66,7 @@ public class TieredIndexImpl implements TieredIndex { * Constructor for building the index using specified dataset */ private TieredIndexImpl( - TieredIndexParams indexParameters, Dataset dataset, CuVSResources resources) { + TieredIndexParams indexParameters, CuVSMatrix dataset, CuVSResources resources) { this.tieredIndexParameters = indexParameters; this.dataset = dataset; this.resources = resources; @@ -97,7 +91,7 @@ private void checkNotDestroyed() { * Invokes the native destroy_tiered_index to de-allocate the Tiered index */ @Override - public void destroyIndex() throws Throwable { + public void destroyIndex() { checkNotDestroyed(); try { int returnValue = cuvsTieredIndexDestroy(tieredIndexReference.getMemorySegment()); @@ -105,7 +99,6 @@ public void destroyIndex() throws Throwable { if (dataset != null) { dataset.close(); } - } finally { destroyed = true; } @@ -121,8 +114,9 @@ public void destroyIndex() throws Throwable { */ private IndexReference build() { try (var localArena = Arena.ofConfined()) { + assert dataset != null; long rows = dataset.size(); - long cols = dataset.dimensions(); + long cols = dataset.columns(); MemorySegment indexParamsMemorySegment = tieredIndexParameters != null @@ -130,7 +124,7 @@ private IndexReference build() { : MemorySegment.NULL; // Get host data - MemorySegment hostDataSeg = ((DatasetImpl) dataset).asMemorySegment(); + MemorySegment hostDataSeg = ((CuVSHostMatrixImpl) dataset).memorySegment(); try (var resourceAccess = resources.access()) { long cuvsRes = resourceAccess.handle(); @@ -145,7 +139,7 @@ private IndexReference build() { // Create tensor from device memory long[] datasetShape = {rows, cols}; MemorySegment datasetTensor = - prepareTensor(localArena, datasetDP, datasetShape, kDLFloat(), 32, 2, kDLCUDA(), 1); + prepareTensor(localArena, datasetDP, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); MemorySegment index = localArena.allocate(cuvsTieredIndex_t); var returnValue = cuvsTieredIndexCreate(index); @@ -219,21 +213,14 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { // Create tensors from device memory long[] queriesShape = {numQueries, vectorDimension}; MemorySegment queriesTensor = - prepareTensor(localArena, queriesDP, queriesShape, kDLFloat(), 32, 2, kDLCUDA(), 1); + prepareTensor(localArena, queriesDP, queriesShape, kDLFloat(), 32, kDLCUDA(), 1); long[] neighborsShape = {numQueries, topK}; MemorySegment neighborsTensor = prepareTensor( - localArena, - neighborsDP, - neighborsShape, - kDLInt(), - 64, - 2, - kDLCUDA(), - 1); // 64-bit int + localArena, neighborsDP, neighborsShape, kDLInt(), 64, kDLCUDA(), 1); // 64-bit int long[] distancesShape = {numQueries, topK}; MemorySegment distancesTensor = - prepareTensor(localArena, distancesDP, distancesShape, kDLFloat(), 32, 2, kDLCUDA(), 1); + prepareTensor(localArena, distancesDP, distancesShape, kDLFloat(), 32, kDLCUDA(), 1); // Sync before prefilter setup returnValue = cuvsStreamSync(cuvsRes); @@ -264,8 +251,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { "cudaMemcpy"); MemorySegment prefilterTensor = - prepareTensor( - localArena, prefilterDP, prefilterShape, kDLUInt(), 32, 1, kDLCUDA(), 1); + prepareTensor(localArena, prefilterDP, prefilterShape, kDLUInt(), 32, kDLCUDA(), 1); cuvsFilter.type(prefilter, 1); // BITSET cuvsFilter.addr(prefilter, prefilterTensor.address()); @@ -330,13 +316,14 @@ public ExtendBuilder extend() { /** * Performs the actual extend operation */ - private void performExtend(Dataset extendDataset) throws Throwable { + private void performExtend(CuVSMatrix extendDataset) { try (var localArena = Arena.ofConfined()) { + assert extendDataset != null; long rows = extendDataset.size(); - long cols = extendDataset.dimensions(); + long cols = extendDataset.columns(); // Get host data - MemorySegment hostDataSeg = ((DatasetImpl) extendDataset).asMemorySegment(); + MemorySegment hostDataSeg = ((CuVSMatrixBaseImpl) extendDataset).memorySegment(); try (var resourceAccess = resources.access()) { long cuvsRes = resourceAccess.handle(); @@ -352,7 +339,7 @@ private void performExtend(Dataset extendDataset) throws Throwable { // Create tensor from device memory long[] datasetShape = {rows, cols}; MemorySegment datasetTensor = - prepareTensor(localArena, datasetDP, datasetShape, kDLFloat(), 32, 2, kDLCUDA(), 1); + prepareTensor(localArena, datasetDP, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); @@ -371,7 +358,7 @@ private void performExtend(Dataset extendDataset) throws Throwable { */ public static class ExtendBuilder implements TieredIndex.ExtendBuilder { private final TieredIndexImpl index; - private Dataset dataset; + private CuVSMatrix dataset; private ExtendBuilder(TieredIndexImpl index) { this.index = index; @@ -379,18 +366,18 @@ private ExtendBuilder(TieredIndexImpl index) { @Override public ExtendBuilder withDataset(float[][] vectors) { - this.dataset = Dataset.ofArray(vectors); + this.dataset = CuVSMatrix.ofArray(vectors); return this; } @Override - public ExtendBuilder withDataset(Dataset dataset) { + public ExtendBuilder withDataset(CuVSMatrix dataset) { this.dataset = dataset; return this; } @Override - public void execute() throws Throwable { + public void execute() { if (dataset == null) { throw new IllegalArgumentException("Must provide a dataset"); } @@ -508,7 +495,7 @@ public static TieredIndex.Builder newBuilder(CuVSResources cuvsResources) { */ public static class Builder implements TieredIndex.Builder { private final CuVSResources resources; - private Dataset dataset; + private CuVSMatrix dataset; private TieredIndexParams params; private TieredIndexType indexType = TieredIndexType.CAGRA; private InputStream inputStream; @@ -531,12 +518,12 @@ public Builder withDataset(float[][] vectors) { if (vectors == null || vectors.length == 0 || vectors[0].length == 0) { throw new IllegalArgumentException("The input vectors cannot be null or empty"); } - this.dataset = Dataset.ofArray(vectors); + this.dataset = CuVSMatrix.ofArray(vectors); return this; } @Override - public Builder withDataset(Dataset dataset) { + public Builder withDataset(CuVSMatrix dataset) { if (this.dataset != null) { throw new IllegalArgumentException("An input dataset can only be specified once"); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index 576b182588..6d08878650 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -264,11 +264,32 @@ public static MemorySegment buildMemorySegment(Arena arena, float[][] data) { long cols = rows > 0 ? data[0].length : 0; MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout(rows * cols, C_FLOAT); MemorySegment dataMemorySegment = arena.allocate(dataMemoryLayout); + copy(dataMemorySegment, data); + return dataMemorySegment; + } + + public static void copy(MemorySegment memorySegment, float[][] data) { + int rows = data.length; + int cols = rows > 0 ? data[0].length : 0; for (int r = 0; r < rows; r++) { - MemorySegment.copy( - data[r], 0, dataMemorySegment, C_FLOAT, (r * cols * C_FLOAT.byteSize()), (int) cols); + MemorySegment.copy(data[r], 0, memorySegment, C_FLOAT, (r * cols * C_FLOAT.byteSize()), cols); + } + } + + public static void copy(MemorySegment memorySegment, int[][] data) { + int rows = data.length; + int cols = rows > 0 ? data[0].length : 0; + for (int r = 0; r < rows; r++) { + MemorySegment.copy(data[r], 0, memorySegment, C_INT, (r * cols * C_INT.byteSize()), cols); + } + } + + public static void copy(MemorySegment memorySegment, byte[][] data) { + int rows = data.length; + int cols = rows > 0 ? data[0].length : 0; + for (int r = 0; r < rows; r++) { + MemorySegment.copy(data[r], 0, memorySegment, C_CHAR, (r * cols * C_CHAR.byteSize()), cols); } - return dataMemorySegment; } public static BitSet concatenate(BitSet[] arr, int maxSizeOfEachBitSet) { @@ -293,7 +314,6 @@ public static BitSet concatenate(BitSet[] arr, int maxSizeOfEachBitSet) { * @param[in] shape the shape of the tensor * @param[in] code the type code of base types * @param[in] bits the shape of the tensor - * @param[in] ndim the number of dimensions * @return DLManagedTensor */ public static MemorySegment prepareTensor( @@ -302,7 +322,6 @@ public static MemorySegment prepareTensor( long[] shape, int code, int bits, - int ndim, int deviceType, int lanes) { @@ -315,6 +334,7 @@ public static MemorySegment prepareTensor( DLDevice.device_type(dlDevice, deviceType); DLTensor.device(dlTensor, dlDevice); + var ndim = shape.length; DLTensor.ndim(dlTensor, ndim); MemorySegment dtype = DLDataType.allocate(arena); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 5940e3db48..668befe236 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -15,24 +15,9 @@ */ package com.nvidia.cuvs.spi; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; - -import com.nvidia.cuvs.BruteForceIndex; -import com.nvidia.cuvs.CagraIndex; -import com.nvidia.cuvs.CagraMergeParams; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.Dataset; -import com.nvidia.cuvs.HnswIndex; -import com.nvidia.cuvs.TieredIndex; -import com.nvidia.cuvs.internal.BruteForceIndexImpl; -import com.nvidia.cuvs.internal.CagraIndexImpl; -import com.nvidia.cuvs.internal.CuVSResourcesImpl; -import com.nvidia.cuvs.internal.DatasetImpl; -import com.nvidia.cuvs.internal.HnswIndexImpl; -import com.nvidia.cuvs.internal.TieredIndexImpl; +import com.nvidia.cuvs.*; +import com.nvidia.cuvs.internal.*; import com.nvidia.cuvs.internal.common.Util; -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; @@ -48,20 +33,26 @@ final class JDKProvider implements CuVSProvider { static MethodHandle createNativeDatasetBuilder() { try { var lookup = MethodHandles.lookup(); - var mt = MethodType.methodType(Dataset.class, MemorySegment.class, int.class, int.class); + var mt = + MethodType.methodType( + CuVSMatrix.class, + MemorySegment.class, + int.class, + int.class, + CuVSMatrix.DataType.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); + private static CuVSMatrix createNativeDataset( + MemorySegment memorySegment, int size, int dimensions, CuVSMatrix.DataType dataType) { + return new CuVSHostMatrixImpl(memorySegment, size, dimensions, dataType); } @Override - public CuVSResources newCuVSResources(Path tempDirectory) throws Throwable { + public CuVSResources newCuVSResources(Path tempDirectory) { Objects.requireNonNull(tempDirectory); if (Files.notExists(tempDirectory)) { throw new IllegalArgumentException("does not exist:" + tempDirectory); @@ -110,46 +101,91 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams merge } @Override - public Dataset.Builder newDatasetBuilder(int size, int dimensions) + public CuVSMatrix.Builder newMatrixBuilder(int size, int dimensions, CuVSMatrix.DataType dataType) throws UnsupportedOperationException { - MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long) size * dimensions, C_FLOAT); - var arena = Arena.ofShared(); - var seg = arena.allocate(dataMemoryLayout); + var dataset = new CuVSHostMatrixArenaImpl(size, dimensions, dataType); - return new Dataset.Builder() { + return new CuVSMatrix.Builder() { int current = 0; @Override public void addVector(float[] vector) { + internalAddVector(vector); + } + + @Override + public void addVector(byte[] vector) { + internalAddVector(vector); + } + + @Override + public void addVector(int[] vector) { + internalAddVector(vector); + } + + private void internalAddVector(Object vector) { if (current >= size) throw new ArrayIndexOutOfBoundsException(); MemorySegment.copy( - vector, 0, seg, C_FLOAT, ((current++) * dimensions * C_FLOAT.byteSize()), dimensions); + vector, + 0, + dataset.memorySegment(), + dataset.valueLayout(), + ((current++) * dimensions * dataset.valueLayout().byteSize()), + dimensions); } @Override - public Dataset build() { - return new DatasetImpl(arena, seg, size, dimensions); + public CuVSMatrix build() { + return dataset; } }; } @Override - public MethodHandle newNativeDatasetBuilder() { + public MethodHandle newNativeMatrixBuilder() { return createNativeDataset$mh; } @Override - public Dataset newArrayDataset(float[][] vectors) { + public CuVSMatrix newMatrixFromArray(float[][] vectors) { + Objects.requireNonNull(vectors); + if (vectors.length == 0) { + throw new IllegalArgumentException("vectors should not be empty"); + } + int size = vectors.length; + int columns = vectors[0].length; + + var dataset = new CuVSHostMatrixArenaImpl(size, columns, CuVSMatrix.DataType.FLOAT); + Util.copy(dataset.memorySegment(), vectors); + return dataset; + } + + @Override + public CuVSMatrix newMatrixFromArray(int[][] vectors) { + Objects.requireNonNull(vectors); + if (vectors.length == 0) { + throw new IllegalArgumentException("vectors should not be empty"); + } + int size = vectors.length; + int columns = vectors[0].length; + + var dataset = new CuVSHostMatrixArenaImpl(size, columns, CuVSMatrix.DataType.INT); + Util.copy(dataset.memorySegment(), vectors); + return dataset; + } + + @Override + public CuVSMatrix newMatrixFromArray(byte[][] vectors) { Objects.requireNonNull(vectors); if (vectors.length == 0) { throw new IllegalArgumentException("vectors should not be empty"); } int size = vectors.length; - int dimensions = vectors[0].length; + int columns = vectors[0].length; - Arena arena = Arena.ofShared(); - var memorySegment = Util.buildMemorySegment(arena, vectors); - return new DatasetImpl(arena, memorySegment, size, dimensions); + var dataset = new CuVSHostMatrixArenaImpl(size, columns, CuVSMatrix.DataType.BYTE); + Util.copy(dataset.memorySegment(), vectors); + return dataset; } } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java index c39c3c0dfa..92f215165b 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java @@ -18,7 +18,8 @@ import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; import java.io.*; -import java.lang.invoke.MethodHandles; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.BitSet; import java.util.List; import java.util.Map; @@ -26,13 +27,9 @@ import java.util.function.LongToIntFunction; import org.junit.Before; import org.junit.Test; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; public class BruteForceAndSearchIT extends CuVSTestCase { - private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); - @Before public void setup() { assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); @@ -67,31 +64,30 @@ private static void indexAndQueryOnce( .build(); // Saving the index on to the disk. - String indexFileName = UUID.randomUUID().toString() + ".bf"; + String indexFileName = UUID.randomUUID() + ".bf"; try (var outputStream = new FileOutputStream(indexFileName)) { index.serialize(outputStream); } // Loading a BRUTEFORCE index from disk. - File indexFile = new File(indexFileName); - InputStream inputStream = new FileInputStream(indexFile); - BruteForceIndex loadedIndex = BruteForceIndex.newBuilder(resources).from(inputStream).build(); - - // search the loaded index - SearchResults results = loadedIndex.search(cuvsQuery); - checkResults(expectedResults, results.getResults()); + Path indexFile = Path.of(indexFileName); + try (var inputStream = Files.newInputStream(indexFile)) { + BruteForceIndex loadedIndex = BruteForceIndex.newBuilder(resources).from(inputStream).build(); - // search the first index - results = index.search(cuvsQuery); - checkResults(expectedResults, results.getResults()); + // search the loaded index + SearchResults results = loadedIndex.search(cuvsQuery); + checkResults(expectedResults, results.getResults()); - // Cleanup - index.destroyIndex(); - loadedIndex.destroyIndex(); + // search the first index + results = index.search(cuvsQuery); + checkResults(expectedResults, results.getResults()); - if (indexFile.exists()) { - indexFile.delete(); + // Cleanup + index.destroyIndex(); + loadedIndex.destroyIndex(); } + + Files.deleteIfExists(indexFile); } /** diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java index a998b3ee9e..fb338106f2 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java @@ -120,7 +120,8 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro BruteForceIndex index; if (useNativeMemoryDataset) { - var datasetBuilder = Dataset.builder(vectors.length, vectors[0].length); + var datasetBuilder = + CuVSMatrix.builder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); for (float[] v : vectors) { datasetBuilder.addVector(v); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index bfbb4e2d1b..10ea10bcf5 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -16,9 +16,7 @@ package com.nvidia.cuvs; import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNotNull; -import static org.junit.Assert.fail; +import static org.junit.Assert.*; import com.carrotsearch.randomizedtesting.RandomizedRunner; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; @@ -139,7 +137,7 @@ public void testIndexingAndSearchingFlow() throws Throwable { int numTestsRuns = 5; try (CuVSResources resources = CheckedCuVSResources.create()) { for (int j = 0; j < numTestsRuns; j++) { - var index = indexOnce(Dataset.ofArray(dataset), resources); + var index = indexOnce(CuVSMatrix.ofArray(dataset), resources); var indexPath = serializeOnce(index); var loadedIndex = deserializeOnce(indexPath, resources); queryAndCompare( @@ -170,7 +168,7 @@ public void testIndexingAndSearchingFlowInDifferentThreads() throws Throwable { runInAnotherThread( () -> { try { - var index = indexOnce(Dataset.ofArray(dataset), resources); + var index = indexOnce(CuVSMatrix.ofArray(dataset), resources); var indexPath = serializeOnce(index); var loadedIndex = deserializeOnce(indexPath, resources); queryAndCompare( @@ -200,12 +198,13 @@ public void testIndexingAndSearchingFlowConcurrently() throws Throwable { List> expectedResults = getExpectedResults(); int numTestsRuns = 10; + runConcurrently( numTestsRuns, () -> () -> { try (CuVSResources resources = CheckedCuVSResources.create()) { - var index = indexOnce(Dataset.ofArray(dataset), resources); + var index = indexOnce(CuVSMatrix.ofArray(dataset), resources); var indexPath = serializeOnce(index); var loadedIndex = deserializeOnce(indexPath, resources); queryAndCompare( @@ -233,7 +232,7 @@ public void testIndexing() throws Throwable { () -> () -> { try (CuVSResources resources = CheckedCuVSResources.create()) { - var index = indexOnce(Dataset.ofArray(dataset), resources); + var index = indexOnce(CuVSMatrix.ofArray(dataset), resources); index.destroyIndex(); } catch (Throwable e) { throw new RuntimeException(e); @@ -252,7 +251,7 @@ public void testSerialization() throws Throwable { () -> () -> { try (CuVSResources resources = CheckedCuVSResources.create()) { - var index = indexOnce(Dataset.ofArray(dataset), resources); + var index = indexOnce(CuVSMatrix.ofArray(dataset), resources); var indexPath = serializeOnce(index); Files.deleteIfExists(indexPath); index.destroyIndex(); @@ -265,7 +264,7 @@ public void testSerialization() throws Throwable { @Test public void testDeserialization() throws Throwable { - var indexPath = createSerializedIndex(Dataset.ofArray(createSampleData())); + var indexPath = createSerializedIndex(CuVSMatrix.ofArray(createSampleData())); for (int i = 0; i < 100; ++i) { int numTestsRuns = 10; runConcurrently( @@ -282,7 +281,7 @@ public void testDeserialization() throws Throwable { Files.deleteIfExists(indexPath); } - private Path createSerializedIndex(Dataset dataset) throws Throwable { + private Path createSerializedIndex(CuVSMatrix dataset) throws Throwable { try (CuVSResources resources = CheckedCuVSResources.create()) { var index = indexOnce(dataset, resources); var indexPath = serializeOnce(index); @@ -293,7 +292,7 @@ private Path createSerializedIndex(Dataset dataset) throws Throwable { @Test public void testIndexingAndSearchingFlowWithCustomMappingFunction() throws Throwable { - var dataset = Dataset.ofArray(createSampleData()); + var dataset = CuVSMatrix.ofArray(createSampleData()); float[][] queries = createSampleQueries(); var expectedResults = List.of( @@ -315,7 +314,7 @@ public void testIndexingAndSearchingFlowWithCustomMappingFunction() throws Throw @Test public void testIndexingAndSearchingFlowWithCustomMappingList() throws Throwable { - var dataset = Dataset.ofArray(createSampleData()); + var dataset = CuVSMatrix.ofArray(createSampleData()); float[][] queries = createSampleQueries(); var mappings = List.of(4, 3, 2, 1); var expectedResults = @@ -404,7 +403,7 @@ public void testPrefilteringReducesResults() throws Throwable { } } - private CagraIndex indexOnce(Dataset dataset, CuVSResources resources) throws Throwable { + private CagraIndex indexOnce(CuVSMatrix dataset, CuVSResources resources) throws Throwable { // Configure index parameters CagraIndexParams indexParams = new CagraIndexParams.Builder() @@ -500,8 +499,10 @@ public void testNativeDatasetEquivalent() throws Throwable { } try (var resources = CuVSResources.create(); - var javaDataset = Dataset.ofArray(sampleData); - var nativeDataset = DatasetHelper.fromMemorySegment(dataMemorySegment, rows, cols)) { + var javaDataset = CuVSMatrix.ofArray(sampleData); + var nativeDataset = + DatasetHelper.fromMemorySegment( + dataMemorySegment, rows, cols, CuVSMatrix.DataType.FLOAT)) { // Indexing with an on-heap and native datasets produce the same results var javaIndex = indexOnce(javaDataset, resources); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java index cb81f8464b..af4965159b 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java @@ -19,7 +19,6 @@ import com.carrotsearch.randomizedtesting.RandomizedRunner; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; -import java.lang.invoke.MethodHandles; import java.util.BitSet; import java.util.List; import org.junit.Before; @@ -31,7 +30,7 @@ @RunWith(RandomizedRunner.class) public class CagraRandomizedIT extends CuVSTestCase { - private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + private static final Logger log = LoggerFactory.getLogger(CagraRandomizedIT.class); @Before public void setup() { @@ -42,7 +41,7 @@ public void setup() { @Test public void testResultsTopKWithRandomValues() throws Throwable { - boolean useNativeMemoryDatasets[] = {true, false}; + boolean[] useNativeMemoryDatasets = {true, false}; for (int i = 0; i < 100; i++) { for (boolean use : useNativeMemoryDatasets) { tmpResultsTopKWithRandomValues(use); @@ -122,7 +121,8 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro CagraIndex index; if (useNativeMemoryDataset) { - var datasetBuilder = Dataset.builder(vectors.length, vectors[0].length); + var datasetBuilder = + CuVSMatrix.builder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); for (float[] v : vectors) { datasetBuilder.addVector(v); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java new file mode 100644 index 0000000000..3518b8acd6 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java @@ -0,0 +1,286 @@ +/* + * 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 static com.carrotsearch.randomizedtesting.RandomizedTest.*; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; + +import com.carrotsearch.randomizedtesting.RandomizedRunner; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; + +@RunWith(RandomizedRunner.class) +public class CuVSMatrixIT extends CuVSTestCase { + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + initializeRandom(); + } + + private static final float DELTA = 1e-9f; + + private static final byte[][] byteData = { + {1, 2, 3}, + {0, 2, 3}, + {4, 1, 3}, + {3, 0, 2}, + {0, 4, 2} + }; + + private int[][] createIntMatrix() { + int rows = randomIntBetween(1, 32); + int cols = randomIntBetween(1, 100); + + int[][] result = new int[rows][cols]; + + for (int r = 0; r < rows; ++r) { + for (int c = 0; c < cols; ++c) { + result[r][c] = randomInt(); + } + } + return result; + } + + private float[][] createFloatMatrix() { + int rows = randomIntBetween(1, 32); + int cols = randomIntBetween(1, 100); + + float[][] result = new float[rows][cols]; + + for (int r = 0; r < rows; ++r) { + for (int c = 0; c < cols; ++c) { + result[r][c] = randomFloat(); + } + } + return result; + } + + @Test + public void testByteDatasetRowGetAccess() { + try (var dataset = CuVSMatrix.ofArray(byteData)) { + for (int n = 0; n < dataset.size(); ++n) { + var row = dataset.getRow(n); + assertEquals(dataset.columns(), row.size()); + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(byteData[n][i], row.getAsByte(i)); + } + } + } + } + + @Test + public void testByteDatasetRowCopy() { + try (var dataset = CuVSMatrix.ofArray(byteData)) { + for (int n = 0; n < dataset.size(); ++n) { + var row = dataset.getRow(n); + assertEquals(dataset.columns(), row.size()); + + var rowCopy = new byte[(int) row.size()]; + row.toArray(rowCopy); + assertArrayEquals(byteData[n], rowCopy); + } + } + } + + @Test + public void testByteDatasetCopy() { + try (var dataset = CuVSMatrix.ofArray(byteData)) { + var dataCopy = new byte[(int) dataset.size()][(int) dataset.columns()]; + dataset.toArray(dataCopy); + for (int n = 0; n < dataset.size(); ++n) { + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(byteData[n][i], dataCopy[n][i]); + } + } + } + } + + @Test + public void testIntDatasetRowGetAccess() { + var intData = createIntMatrix(); + try (var dataset = CuVSMatrix.ofArray(intData)) { + for (int n = 0; n < dataset.size(); ++n) { + var row = dataset.getRow(n); + assertEquals(dataset.columns(), row.size()); + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(intData[n][i], row.getAsInt(i)); + } + } + } + } + + @Test + public void testIntDatasetRowCopy() { + var intData = createIntMatrix(); + try (var dataset = CuVSMatrix.ofArray(intData)) { + for (int n = 0; n < dataset.size(); ++n) { + var row = dataset.getRow(n); + assertEquals(dataset.columns(), row.size()); + + var rowCopy = new int[(int) row.size()]; + row.toArray(rowCopy); + assertArrayEquals(intData[n], rowCopy); + } + } + } + + @Test + public void testIntDatasetCopy() { + var intData = createIntMatrix(); + try (var dataset = CuVSMatrix.ofArray(intData)) { + var intDataCopy = new int[(int) dataset.size()][(int) dataset.columns()]; + dataset.toArray(intDataCopy); + for (int n = 0; n < dataset.size(); ++n) { + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(intData[n][i], intDataCopy[n][i]); + } + } + } + } + + @Test + public void testFloatDatasetRowGetAccess() { + var floatData = createFloatMatrix(); + try (var dataset = CuVSMatrix.ofArray(floatData)) { + for (int n = 0; n < dataset.size(); ++n) { + var row = dataset.getRow(n); + assertEquals(dataset.columns(), row.size()); + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(floatData[n][i], row.getAsFloat(i), DELTA); + } + } + } + } + + @Test + public void testFloatDatasetRowCopy() { + var floatData = createFloatMatrix(); + try (var dataset = CuVSMatrix.ofArray(floatData)) { + for (int n = 0; n < dataset.size(); ++n) { + var row = dataset.getRow(n); + assertEquals(dataset.columns(), row.size()); + + var rowCopy = new float[(int) row.size()]; + row.toArray(rowCopy); + assertArrayEquals(floatData[n], rowCopy, DELTA); + } + } + } + + @Test + public void testFloatDatasetCopy() { + var floatData = createFloatMatrix(); + try (var dataset = CuVSMatrix.ofArray(floatData)) { + var dataCopy = new float[(int) dataset.size()][(int) dataset.columns()]; + dataset.toArray(dataCopy); + for (int n = 0; n < dataset.size(); ++n) { + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(floatData[n][i], dataCopy[n][i], DELTA); + } + } + } + } + + public void testFloatDatasetBuilder() { + int rows = randomIntBetween(1, 32); + int cols = randomIntBetween(1, 100); + + float[][] data = new float[rows][cols]; + for (int r = 0; r < rows; ++r) { + for (int c = 0; c < cols; ++c) { + data[c][r] = randomFloat(); + } + } + + var builder = CuVSMatrix.builder(rows, cols, CuVSMatrix.DataType.FLOAT); + for (int r = 0; r < rows; ++r) { + builder.addVector(data[r]); + } + + float[][] roundTripData = new float[rows][cols]; + + try (var dataset = builder.build()) { + dataset.toArray(roundTripData); + + for (int n = 0; n < dataset.size(); ++n) { + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(data[n][i], roundTripData[n][i], DELTA); + } + } + } + } + + public void testIntDatasetBuilder() { + int rows = randomIntBetween(1, 32); + int cols = randomIntBetween(1, 100); + + var data = new int[rows][cols]; + for (int r = 0; r < rows; ++r) { + for (int c = 0; c < cols; ++c) { + data[c][r] = randomInt(); + } + } + + var builder = CuVSMatrix.builder(rows, cols, CuVSMatrix.DataType.INT); + for (int r = 0; r < rows; ++r) { + builder.addVector(data[r]); + } + + var roundTripData = new int[rows][cols]; + + try (var dataset = builder.build()) { + dataset.toArray(roundTripData); + + for (int n = 0; n < dataset.size(); ++n) { + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(data[n][i], roundTripData[n][i]); + } + } + } + } + + public void testByteDatasetBuilder() { + int rows = randomIntBetween(1, 32); + int cols = randomIntBetween(1, 100); + + var data = new byte[rows][cols]; + for (int r = 0; r < rows; ++r) { + for (int c = 0; c < cols; ++c) { + data[c][r] = randomByte(); + } + } + + var builder = CuVSMatrix.builder(rows, cols, CuVSMatrix.DataType.BYTE); + for (int r = 0; r < rows; ++r) { + builder.addVector(data[r]); + } + + var roundTripData = new byte[rows][cols]; + + try (var dataset = builder.build()) { + dataset.toArray(roundTripData); + + for (int n = 0; n < dataset.size(); ++n) { + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(data[n][i], roundTripData[n][i]); + } + } + } + } +} diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/DatasetHelper.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/DatasetHelper.java index ae62e10c34..9138f9952a 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/DatasetHelper.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/DatasetHelper.java @@ -22,11 +22,12 @@ public class DatasetHelper { private static final MethodHandle createDataset$mh = - CuVSProvider.provider().newNativeDatasetBuilder(); + CuVSProvider.provider().newNativeMatrixBuilder(); - public static Dataset fromMemorySegment(MemorySegment memorySegment, int size, int dimensions) { + public static CuVSMatrix fromMemorySegment( + MemorySegment memorySegment, int size, int dimensions, CuVSMatrix.DataType dataType) { try { - return (Dataset) createDataset$mh.invokeExact(memorySegment, size, dimensions); + return (CuVSMatrix) createDataset$mh.invokeExact(memorySegment, size, dimensions, dataType); } catch (Throwable e) { if (e instanceof Error err) { throw err; diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java index d2e198738b..ded5bda54f 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java @@ -113,7 +113,8 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro // Create the index with the dataset final CagraIndex index; if (useNativeMemoryDataset) { - var datasetBuilder = Dataset.builder(vectors.length, vectors[0].length); + var datasetBuilder = + CuVSMatrix.builder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); for (float[] v : vectors) { datasetBuilder.addVector(v); }