From 1155a3a427cd1d1bfaf8fe74a937ed6dfa797ec7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Sat, 26 Jul 2025 03:58:49 +0200 Subject: [PATCH] [Java] Extend `Dataset` to work as an output data container (#1111) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In #902 and https://github.com/rapidsai/cuvs/pull/1034 we introduced a `Dataset` interface to support on-heap and off-heap ("native") memory seamlessly as inputs for cagra and bruteforce index building. As we expand the functionality of cuvs-java, we realized we have similar needs for outputs (see e.g. https://github.com/rapidsai/cuvs/pull/1105 / https://github.com/rapidsai/cuvs/pull/1102 or https://github.com/rapidsai/cuvs/pull/1104). This PR extends `Dataset` to support being used as an output, wrapping native (off-heap) memory in a convenient and efficient way, and providing common utilities to transform to and from on-heap memory. This work is inspired by the existing raft `mdspan` and `DLTensor` data structures, but tailored to our needs (2d only, just 3 data types, etc.). The PR keeps the current implementation simple and minimal on purpose, but structured in a way that is simple to extend. By itself, the PR is just a refactoring to extend the `Dataset` implementation and reorganize the implementation classes; its real usefulness will be in using it in the PRs mentioned above (in fact, this PR has been extracted from https://github.com/rapidsai/cuvs/pull/1105). The implementation class hierarchy is implemented with future extensions in mind: atm we have one `HostMemoryDatasetImpl`, but we are already thinking to have a corresponding `DeviceMemoryDatasetImpl` that will wrap and manage (views) on GPU memory to avoid (in some cases) extra copies of data from GPU memory to CPU memory only to process them or forward them to another algorithm (e.g quantization followed by indexing). Future work will also include add support/refactoring to allocate and manage GPU memory and DLTensors (e.g. working better with/refactoring `prepareTensor`). Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1111 --- .../com/nvidia/cuvs/CagraIndexBenchmarks.java | 10 +- .../java/com/nvidia/cuvs/BruteForceIndex.java | 4 +- .../main/java/com/nvidia/cuvs/CagraIndex.java | 4 +- .../com/nvidia/cuvs/CuVSDeviceMatrix.java | 21 ++ .../java/com/nvidia/cuvs/CuVSHostMatrix.java | 23 ++ .../main/java/com/nvidia/cuvs/CuVSMatrix.java | 156 ++++++++++ .../main/java/com/nvidia/cuvs/Dataset.java | 74 ----- .../main/java/com/nvidia/cuvs/RowView.java | 74 +++++ .../java/com/nvidia/cuvs/TieredIndex.java | 8 +- .../com/nvidia/cuvs/spi/CuVSProvider.java | 29 +- .../nvidia/cuvs/spi/UnsupportedProvider.java | 19 +- .../cuvs/internal/BruteForceIndexImpl.java | 33 +- .../nvidia/cuvs/internal/CagraIndexImpl.java | 35 ++- ...Impl.java => CuVSHostMatrixArenaImpl.java} | 49 +-- .../cuvs/internal/CuVSHostMatrixImpl.java | 198 ++++++++++++ .../cuvs/internal/CuVSMatrixBaseImpl.java | 48 +++ .../nvidia/cuvs/internal/HnswIndexImpl.java | 7 +- .../nvidia/cuvs/internal/TieredIndexImpl.java | 63 ++-- .../com/nvidia/cuvs/internal/common/Util.java | 30 +- .../com/nvidia/cuvs/spi/JDKProvider.java | 108 ++++--- .../nvidia/cuvs/BruteForceAndSearchIT.java | 38 ++- .../nvidia/cuvs/BruteForceRandomizedIT.java | 3 +- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 31 +- .../com/nvidia/cuvs/CagraRandomizedIT.java | 8 +- .../java/com/nvidia/cuvs/CuVSMatrixIT.java | 286 ++++++++++++++++++ .../java/com/nvidia/cuvs/DatasetHelper.java | 7 +- .../com/nvidia/cuvs/HnswRandomizedIT.java | 3 +- 27 files changed, 1081 insertions(+), 288 deletions(-) create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java delete mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/Dataset.java create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/RowView.java rename java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/{DatasetImpl.java => CuVSHostMatrixArenaImpl.java} (54%) create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java 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); }