diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java index 6dcc4c892f..98f269b573 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java @@ -79,6 +79,7 @@ private boolean mustClose() { public void close() { if (mustClose()) { checkCuVSError(cuvsRMMFree(cuvsResourceHandle, pointer, numBytes), "cuvsRMMFree"); + pointer = MemorySegment.NULL; } } 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 c0c0c8a37a..45076e6482 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,15 +16,13 @@ package com.nvidia.cuvs; import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static com.nvidia.cuvs.CuVSMatrixIT.assertSame2dArray; import static org.junit.Assert.*; import com.carrotsearch.randomizedtesting.RandomizedRunner; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; import com.nvidia.cuvs.CagraMergeParams.MergeStrategy; -import java.io.File; -import java.io.FileInputStream; -import java.io.FileOutputStream; import java.io.InputStream; import java.lang.foreign.Arena; import java.lang.foreign.Linker; @@ -397,12 +395,12 @@ public void testPrefilteringReducesResults() throws Throwable { .withMetric(CuvsDistanceType.L2Expanded) .build(); - try (CuVSResources resources = CheckedCuVSResources.create()) { - CagraIndex index = - CagraIndex.newBuilder(resources) - .withDataset(dataset) - .withIndexParams(indexParams) - .build(); + try (CuVSResources resources = CheckedCuVSResources.create(); + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build()) { // No prefilter (all points allowed) CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); @@ -514,6 +512,10 @@ private void cleanup(CagraIndex index, CagraIndex loadedIndex) throws Throwable loadedIndex.close(); } + /** + * Tests that an index built starting from a native MemorySegment is identical to one built from + * Java heap arrays + */ @Test public void testNativeDatasetEquivalent() throws Throwable { float[][] sampleData = createSampleData(); @@ -555,6 +557,37 @@ public void testNativeDatasetEquivalent() throws Throwable { } } + /** + * Tests that an index built starting from device memory ({@link CuVSDeviceMatrix}) is identical to one + * built from Java heap arrays + */ + @Test + public void testDeviceDatasetEquivalent() throws Throwable { + float[][] sampleData = createSampleData(); + + try (var resources = CuVSResources.create(); + var javaDataset = CuVSMatrix.ofArray(sampleData); + var deviceDataset = javaDataset.toDevice(resources)) { + + // Indexing with an on-heap and native datasets produce the same results + var javaIndex = indexOnce(javaDataset, resources); + var deviceIndex = indexOnce(deviceDataset, resources); + + int size = (int) javaIndex.getGraph().size(); + assertEquals(size, (int) deviceIndex.getGraph().size()); + + int columns = (int) javaIndex.getGraph().columns(); + assertEquals(columns, (int) deviceIndex.getGraph().columns()); + + var javaIndexGraph = new int[size][columns]; + var deviceIndexGraph = new int[size][columns]; + javaIndex.getGraph().toArray(javaIndexGraph); + deviceIndex.getGraph().toArray(deviceIndexGraph); + + assertSame2dArray(size, columns, javaIndexGraph, deviceIndexGraph); + } + } + @Test public void testMergingIndexes() throws Throwable { float[][] vector1 = { @@ -626,22 +659,23 @@ public void testMergingIndexes() throws Throwable { // --- Serialization/deserialization check --- String indexFileName = UUID.randomUUID() + ".cag"; - mergedIndex.serialize(new FileOutputStream(indexFileName)); + var indexFile = Path.of(indexFileName); - File indexFile = new File(indexFileName); - InputStream inputStream = new FileInputStream(indexFile); - CagraIndex loadedMergedIndex = CagraIndex.newBuilder(resources).from(inputStream).build(); + try (var out = Files.newOutputStream(indexFile)) { + mergedIndex.serialize(out); + } - SearchResults resultsFromLoaded = loadedMergedIndex.search(query); - assertEquals(expectedResults, resultsFromLoaded.getResults()); + try (InputStream inputStream = Files.newInputStream(indexFile)) { + CagraIndex loadedMergedIndex = CagraIndex.newBuilder(resources).from(inputStream).build(); - if (indexFile.exists()) { - indexFile.delete(); + SearchResults resultsFromLoaded = loadedMergedIndex.search(query); + assertEquals(expectedResults, resultsFromLoaded.getResults()); + mergedIndex.close(); + loadedMergedIndex.close(); } + Files.deleteIfExists(indexFile); index1.close(); index2.close(); - mergedIndex.close(); - loadedMergedIndex.close(); } } @@ -710,50 +744,51 @@ public void testMergeStrategies() throws Throwable { .build(); log.trace("Merging indexes with PHYSICAL strategy..."); - CagraIndex physicalMergedIndex = - CagraIndex.merge(new CagraIndex[] {index1, index2}, physicalMergeParams); - log.trace("Physical merge completed successfully"); + try (CagraIndex physicalMergedIndex = + CagraIndex.merge(new CagraIndex[] {index1, index2}, physicalMergeParams)) { + log.trace("Physical merge completed successfully"); + + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(3) + .withSearchParams(searchParams) + .withQueryVectors(queries) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + + log.trace("Searching physically merged index..."); + SearchResults physicalResults = physicalMergedIndex.search(query); + assertNotNull("Physical merge search results should not be null", physicalResults); + assertEquals( + "Physical merge search results should match expected", + expectedResults, + physicalResults.getResults()); - CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + // --- Serialization/deserialization check for both merged indexes --- + String physicalIndexFileName = UUID.randomUUID() + ".cag"; + var physicalIndexFile = Path.of(physicalIndexFileName); - CagraQuery query = - new CagraQuery.Builder(resources) - .withTopK(3) - .withSearchParams(searchParams) - .withQueryVectors(queries) - .withMapping(SearchResults.IDENTITY_MAPPING) - .build(); + try (var out = Files.newOutputStream(physicalIndexFile)) { + physicalMergedIndex.serialize(out); + } - log.trace("Searching physically merged index..."); - SearchResults physicalResults = physicalMergedIndex.search(query); - assertNotNull("Physical merge search results should not be null", physicalResults); - assertEquals( - "Physical merge search results should match expected", - expectedResults, - physicalResults.getResults()); - - // --- Serialization/deserialization check for both merged indexes --- - String physicalIndexFileName = UUID.randomUUID().toString() + ".cag"; - physicalMergedIndex.serialize(new FileOutputStream(physicalIndexFileName)); - - File physicalIndexFile = new File(physicalIndexFileName); - InputStream physicalInputStream = new FileInputStream(physicalIndexFile); - CagraIndex loadedPhysicalIndex = - CagraIndex.newBuilder(resources).from(physicalInputStream).build(); - - SearchResults resultsFromLoadedPhysical = loadedPhysicalIndex.search(query); - assertEquals( - "Loaded physical index search results should match expected", - expectedResults, - resultsFromLoadedPhysical.getResults()); - - if (physicalIndexFile.exists()) { - physicalIndexFile.delete(); + try (InputStream physicalInputStream = Files.newInputStream(physicalIndexFile); + CagraIndex loadedPhysicalIndex = + CagraIndex.newBuilder(resources).from(physicalInputStream).build()) { + + Files.deleteIfExists(physicalIndexFile); + + SearchResults resultsFromLoadedPhysical = loadedPhysicalIndex.search(query); + assertEquals( + "Loaded physical index search results should match expected", + expectedResults, + resultsFromLoadedPhysical.getResults()); + } } index1.close(); index2.close(); - physicalMergedIndex.close(); - loadedPhysicalIndex.close(); } } } 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 6e02793478..d7d41c8c2f 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 @@ -39,17 +39,26 @@ public void setup() { log.trace("Random context initialized for test."); } + enum TestDatasetMemoryKind { + HEAP, + NATIVE, + DEVICE + } + @Test public void testResultsTopKWithRandomValues() throws Throwable { - boolean[] useNativeMemoryDatasets = {true, false}; + TestDatasetMemoryKind[] testDatasetMemoryKinds = { + TestDatasetMemoryKind.HEAP, TestDatasetMemoryKind.NATIVE, TestDatasetMemoryKind.DEVICE + }; for (int i = 0; i < 100; i++) { - for (boolean use : useNativeMemoryDatasets) { - tmpResultsTopKWithRandomValues(use); + for (var datasetMemoryKind : testDatasetMemoryKinds) { + tmpResultsTopKWithRandomValues(datasetMemoryKind); } } } - private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) throws Throwable { + private void tmpResultsTopKWithRandomValues(TestDatasetMemoryKind datasetMemoryKind) + throws Throwable { int DATASET_SIZE_LIMIT = 10_000; int DIMENSIONS_LIMIT = 2048; int NUM_QUERIES_LIMIT = 10; @@ -90,7 +99,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro log.debug("Dataset size: {}x{}", datasetSize, dimensions); log.debug("Query size: {}x{}", numQueries, dimensions); log.debug("TopK: {}", topK); - log.debug("Use native memory dataset? " + useNativeMemoryDataset); + log.debug("Use memory dataset: " + datasetMemoryKind.name()); // Debugging: Log dataset and queries if (log.isDebugEnabled()) { @@ -119,8 +128,8 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) .build(); - CagraIndex index; - if (useNativeMemoryDataset) { + final CagraIndex index; + if (datasetMemoryKind == TestDatasetMemoryKind.NATIVE) { var datasetBuilder = CuVSMatrix.hostBuilder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); for (float[] v : vectors) { @@ -131,7 +140,20 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro .withDataset(datasetBuilder.build()) .withIndexParams(indexParams) .build(); + } else if (datasetMemoryKind == TestDatasetMemoryKind.DEVICE) { + var datasetBuilder = + CuVSMatrix.deviceBuilder( + resources, vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); + for (float[] v : vectors) { + datasetBuilder.addVector(v); + } + index = + CagraIndex.newBuilder(resources) + .withDataset(datasetBuilder.build()) + .withIndexParams(indexParams) + .build(); } else { + assert datasetMemoryKind == TestDatasetMemoryKind.HEAP; index = CagraIndex.newBuilder(resources) .withDataset(vectors) 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 index 499c7d366f..917acd4fc4 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java @@ -211,11 +211,7 @@ public void testIntDatasetCopy() { 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]); - } - } + assertSame2dArray(dataset.size(), dataset.columns(), intData, intDataCopy); } } @@ -254,10 +250,32 @@ public void testFloatDatasetCopy() { 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); - } + assertSame2dArray(dataset.size(), dataset.columns(), floatData, dataCopy); + } + } + + static void assertSame2dArray(long rows, long cols, float[][] array1, float[][] array2) { + assertEquals(rows, array1.length); + assertEquals(cols, array1[0].length); + assertEquals(rows, array2.length); + assertEquals(cols, array2[0].length); + + for (int n = 0; n < rows; ++n) { + for (int i = 0; i < cols; ++i) { + assertEquals(array1[n][i], array2[n][i], DELTA); + } + } + } + + static void assertSame2dArray(long rows, long cols, int[][] array1, int[][] array2) { + assertEquals(rows, array1.length); + assertEquals(cols, array1[0].length); + assertEquals(rows, array2.length); + assertEquals(cols, array2[0].length); + + for (int n = 0; n < rows; ++n) { + for (int i = 0; i < cols; ++i) { + assertEquals(array1[n][i], array2[n][i]); } } } @@ -280,11 +298,7 @@ private void testFloatDatasetBuilder(int rows, int cols, CuVSMatrix.Builder b 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); - } - } + assertSame2dArray(dataset.size(), dataset.columns(), data, roundTripData); } } @@ -323,12 +337,7 @@ private void testIntDatasetBuilder(int rows, int cols, CuVSMatrix.Builder bui 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]); - } - } + assertSame2dArray(dataset.size(), dataset.columns(), data, roundTripData); } } @@ -421,11 +430,7 @@ public void testDeviceToHost() throws Throwable { hostMatrix.toArray(roundTripData); - for (int n = 0; n < hostMatrix.size(); ++n) { - for (int i = 0; i < hostMatrix.columns(); ++i) { - assertEquals(data[n][i], roundTripData[n][i], 1e-9); - } - } + assertSame2dArray(hostMatrix.size(), hostMatrix.columns(), data, roundTripData); } } } @@ -452,11 +457,7 @@ public void testHostToDevice() throws Throwable { deviceMatrix.toArray(roundTripData); - for (int n = 0; n < deviceMatrix.size(); ++n) { - for (int i = 0; i < deviceMatrix.columns(); ++i) { - assertEquals(data[n][i], roundTripData[n][i], 1e-9); - } - } + assertSame2dArray(deviceMatrix.size(), deviceMatrix.columns(), data, roundTripData); } } } @@ -481,11 +482,7 @@ public void testHostToHostReturnsWeakReferenceSameData() { hostMatrix2.toArray(roundTripData); - for (int n = 0; n < hostMatrix2.size(); ++n) { - for (int i = 0; i < hostMatrix2.columns(); ++i) { - assertEquals(data[n][i], roundTripData[n][i], 1e-9); - } - } + assertSame2dArray(hostMatrix2.size(), hostMatrix2.columns(), data, roundTripData); } } @@ -517,11 +514,7 @@ public void testDeviceToDeviceReturnsWeakReferenceSameData() throws Throwable { deviceMatrix2.toArray(roundTripData); - for (int n = 0; n < deviceMatrix2.size(); ++n) { - for (int i = 0; i < deviceMatrix2.columns(); ++i) { - assertEquals(data[n][i], roundTripData[n][i], 1e-9); - } - } + assertSame2dArray(deviceMatrix2.size(), deviceMatrix2.columns(), data, roundTripData); } } }