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 52a8714448..79cf23b4f3 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 @@ -6,6 +6,8 @@ import static com.nvidia.cuvs.internal.CuVSParamsHelper.*; import static com.nvidia.cuvs.internal.common.CloseableRMMAllocation.allocateRMMSegment; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT_BYTE_SIZE; import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.DEVICE_TO_HOST; @@ -224,13 +226,12 @@ public SearchResults search(CagraQuery query) throws Throwable { long numBlocks = topK * numQueries; SequenceLayout neighborsSequenceLayout = MemoryLayout.sequenceLayout(numBlocks, C_INT); - SequenceLayout distancesSequenceLayout = - MemoryLayout.sequenceLayout(numBlocks, queryVectors.valueLayout()); + SequenceLayout distancesSequenceLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); MemorySegment neighborsMemorySegment = localArena.allocate(neighborsSequenceLayout); MemorySegment distancesMemorySegment = localArena.allocate(distancesSequenceLayout); final long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; - final long distancesBytes = queryVectors.valueLayout().byteSize() * numQueries * topK; + final long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; final boolean hasPreFilter = query.getPrefilter() != null; final BitSet[] prefilters = hasPreFilter ? new BitSet[] {query.getPrefilter()} : EMPTY_PREFILTER_BITSET; @@ -259,12 +260,7 @@ public SearchResults search(CagraQuery query) throws Throwable { long[] distancesShape = {numQueries, topK}; MemorySegment distancesTensor = prepareTensor( - localArena, - distancesDP.handle(), - distancesShape, - deviceQueryVectors.code(), - deviceQueryVectors.bits(), - kDLCUDA()); + localArena, distancesDP.handle(), distancesShape, kDLFloat(), 32, kDLCUDA()); // prepare the prefiltering data MemorySegment prefilter = cuvsFilter.allocate(localArena); @@ -633,8 +629,10 @@ private static void populateNativeIndexParams( cuvsAceParams.npartitions(cuvsAceParamsMemorySegment, cuVSAceParams.getNpartitions()); cuvsAceParams.ef_construction(cuvsAceParamsMemorySegment, cuVSAceParams.getEfConstruction()); cuvsAceParams.use_disk(cuvsAceParamsMemorySegment, cuVSAceParams.isUseDisk()); - cuvsAceParams.max_host_memory_gb(cuvsAceParamsMemorySegment, cuVSAceParams.getMaxHostMemoryGb()); - cuvsAceParams.max_gpu_memory_gb(cuvsAceParamsMemorySegment, cuVSAceParams.getMaxGpuMemoryGb()); + cuvsAceParams.max_host_memory_gb( + cuvsAceParamsMemorySegment, cuVSAceParams.getMaxHostMemoryGb()); + cuvsAceParams.max_gpu_memory_gb( + cuvsAceParamsMemorySegment, cuVSAceParams.getMaxGpuMemoryGb()); String buildDir = cuVSAceParams.getBuildDir(); if (buildDir != null && !buildDir.isEmpty()) { 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 c5216b6086..24a655740b 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 @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -595,6 +595,50 @@ public void testPrefilteringReducesResults() throws Throwable { } } + /** + * Regression test for a bug in {@code CagraIndexImpl.search}: the {@code distances} output tensor + * was described (dtype and buffer size) using the query vectors' data type instead of the + * float32 type the C API mandates. Float queries masked the bug because their dtype already matches + * float32. With byte (uint8) queries the distances tensor became uint8/8-bit and its device buffer + * was sized at 1 byte per value instead of 4, which the C wrapper rejects up front + * ("distances should be of type float32") -- and would otherwise be a GPU buffer overflow. The + * existing byte tests only cover build/serialize/deserialize, never search, so this path was + * uncovered. + */ + @Test + public void testByteQuerySearch() throws Throwable { + // Small, unambiguous byte dataset (values within signed-byte range, mapped to uint8). + byte[][] dataset = { + {0, 0}, + {5, 5}, + {50, 50}, + {100, 100} + }; + // Each query equals a dataset row, so its nearest neighbor is that row at distance 0. + byte[][] queries = { + {0, 0}, // -> id 0 + {100, 100} // -> id 3 + }; + List> expectedResults = List.of(Map.of(0, 0.0f), Map.of(3, 0.0f)); + + try (CuVSResources resources = CheckedCuVSResources.create(); + var index = indexOnce(CuVSMatrix.ofArray(dataset), resources); + var queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(1) + .withSearchParams(new CagraSearchParams.Builder().build()) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + + // Fails with "distances should be of type float32" when the bug is present. + SearchResults results = index.search(query); + log.debug("Byte-query search results: {}", results.getResults()); + checkResults(expectedResults, results.getResults()); + } + } + private CagraIndex indexOnce(CuVSMatrix dataset, CuVSResources resources) throws Throwable { // Configure index parameters CagraIndexParams indexParams =