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 ea73d43b69..a494b72e36 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 @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.internal; @@ -184,10 +184,13 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { final MemorySegment prefilterDataMemorySegment; BitSet[] prefilters = cuvsQuery.getPrefilters(); if (prefilters != null && prefilters.length > 0) { + prefilterDataLength = (long) cuvsQuery.getNumDocs() * prefilters.length; BitSet concatenatedFilters = concatenate(prefilters, cuvsQuery.getNumDocs()); + // Size the host buffer by the full logical bit count; BitSet.toLongArray() would drop the + // trailing all-zero words, leaving the prefilterBytes copy reading past the buffer's end. long[] filters = concatenatedFilters.toLongArray(); - prefilterDataMemorySegment = buildMemorySegment(localArena, filters); - prefilterDataLength = (long) cuvsQuery.getNumDocs() * prefilters.length; + prefilterDataMemorySegment = + buildMemorySegment(localArena, filters, (prefilterDataLength + 63) / 64); long[] prefilterShape = {(prefilterDataLength + 31) / 32}; prefilterBytes = C_INT_BYTE_SIZE * prefilterShape[0]; } else { 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 abc53a5945..52a8714448 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 @@ -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.internal; @@ -275,7 +275,8 @@ public SearchResults search(CagraQuery query) throws Throwable { } else { BitSet concatenatedFilters = concatenate(prefilters, query.getNumDocs()); long[] filters = concatenatedFilters.toLongArray(); - var prefilterDataMemorySegment = buildMemorySegment(localArena, filters); + var prefilterDataMemorySegment = + buildMemorySegment(localArena, filters, (prefilterDataLength + 63) / 64); long[] prefilterShape = {prefilterLen}; 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 cfd3a1aa1a..148c2b4f91 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 @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.internal; @@ -244,7 +244,8 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { if (hasPreFilter) { BitSet concatenatedFilters = concatenate(prefilters, (int) query.getNumDocs()); long[] filters = concatenatedFilters.toLongArray(); - MemorySegment hostPrefilterSeg = buildMemorySegment(localArena, filters); + MemorySegment hostPrefilterSeg = + buildMemorySegment(localArena, filters, (prefilterDataLength + 63) / 64); long[] prefilterShape = {prefilterLen}; 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 7ca7640cd8..fa2227041d 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 @@ -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.internal.common; @@ -9,6 +9,11 @@ import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG; import static com.nvidia.cuvs.internal.panama.headers_h.*; +import static com.nvidia.cuvs.internal.panama.headers_h_1.C_POINTER; +import static com.nvidia.cuvs.internal.panama.headers_h_1.cudaMemcpyDeviceToDevice; +import static com.nvidia.cuvs.internal.panama.headers_h_1.cudaMemcpyDeviceToHost; +import static com.nvidia.cuvs.internal.panama.headers_h_1.cudaMemcpyHostToDevice; +import static com.nvidia.cuvs.internal.panama.headers_h_1.cudaMemcpyHostToHost; import static com.nvidia.cuvs.internal.panama.headers_h_1.cudaStream_t; import com.nvidia.cuvs.CuVSResources; @@ -241,10 +246,20 @@ public static MemorySegment duplicateNativeString(String str) { * @return an instance of {@link MemorySegment} */ public static MemorySegment buildMemorySegment(Arena arena, long[] data) { - int cells = data.length; - MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout(cells, C_LONG); + return buildMemorySegment(arena, data, data.length); + } + + /** + * A utility method for building a {@link MemorySegment} for a 1D long array of given size. + * + * @param data The 1D long array for which the {@link MemorySegment} is needed + * @param lengthInLongs if lengthInLongs is longer then the data array, 0s will be appended to the end of the array + * @return an instance of {@link MemorySegment} + */ + public static MemorySegment buildMemorySegment(Arena arena, long[] data, long lengthInLongs) { + MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout(lengthInLongs, C_LONG); MemorySegment dataMemorySegment = arena.allocate(dataMemoryLayout); - MemorySegment.copy(data, 0, dataMemorySegment, C_LONG, 0, cells); + MemorySegment.copy(data, 0, dataMemorySegment, C_LONG, 0, data.length); return dataMemorySegment; } 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 08b26c4b66..d348c44c3f 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 @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -140,6 +140,55 @@ public void testIndexingAndSearchingWithFiltering() throws Throwable { } } + /** + * Test for the prefilter buffer under-allocation. + */ + @Test + public void testPrefilterWithTrailingZeroWords() throws Throwable { + // 128 documents => the concatenated bitmap spans multiple 64-bit words. Each document i sits at + // (i, 0), so document i is at squared-L2 distance (i - 64)^2 from the query below. + int numDocs = 128; + float[][] largeDataset = new float[numDocs][2]; + for (int i = 0; i < numDocs; i++) { + largeDataset[i] = new float[] {i, 0.0f}; + } + // Query sits on the (excluded) nearest document id 64; the only allowed document, id 0, is the + // farthest of all - so any leaked document from the corrupted high words would displace it. + float[][] localQueries = {{64.0f, 0.0f}}; + + // Allow only document 0. All higher bits are 0, so BitSet.toLongArray() drops the trailing words + // covering documents 64..127 - exactly the region the buggy copy fills from adjacent host memory. + BitSet prefilter = new BitSet(numDocs); + prefilter.set(0); + + // id 0 is at squared-L2 distance 64^2 = 4096 from the query. + List> expected = List.of(Map.of(0, 4096.0f)); + + BruteForceIndexParams indexParams = + new BruteForceIndexParams.Builder().withNumWriterThreads(32).build(); + + for (int j = 0; j < 20; j++) { + try (CuVSResources resources = CheckedCuVSResources.create(); + BruteForceIndex index = + BruteForceIndex.newBuilder(resources) + .withDataset(largeDataset) + .withIndexParams(indexParams) + .build()) { + BruteForceQuery cuvsQuery = + new BruteForceQuery.Builder(resources) + .withTopK(1) + .withQueryVectors(localQueries) + .withPrefilters(new BitSet[] {prefilter}, numDocs) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + + // With the bug, an excluded high-index document (closer than id 0) can appear here. + SearchResults results = index.search(cuvsQuery); + checkResults(expected, results.getResults()); + } + } + } + @Test public void testIndexingAndSearchingWithFunctionMapping() throws Throwable { // Expected search results