Forward-merge release/26.08 into main

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 {

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};
 

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};
 

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;
   }
 

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<Map<Integer, Float>> 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