Forward-merge release/26.08 into main

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()) {

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 <em>query</em> 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<Map<Integer, Float>> 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 =