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