From b407df8d71bd61fe3d6d62f7856a8fbc84c131e4 Mon Sep 17 00:00:00 2001
From: nvzm123 <zmeeks@nvidia.com>
Date: Tue, 6 Oct 2026 05:11:16 +0000
Subject: [PATCH] Close CAGRA query matrices after search

---
 .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 28 ++++++-------------
 1 file changed, 9 insertions(+), 19 deletions(-)

diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java
index 1efa2d1174..58ac0d622c 100644
--- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java
+++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java
@@ -526,33 +526,23 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
           // Setting itopK as topK because in any case iTopK should be ATLEAST equal to topK
           searchParams = new CagraSearchParams.Builder().withItopkSize(topK).build();
         }
-        CagraIndex cagraIndex = cuvsIndex.getCagraIndex();
-        assert cagraIndex != null;
-        CagraQuery query = null;
-
         CuVSMatrix.Builder<?> builder =
             CuVSMatrix.deviceBuilder(
                 getCuVSResourcesInstance(), 1, target.length, CuVSMatrix.DataType.FLOAT);
         builder.addVector(target);
-        CuVSMatrix queryVector = builder.build();
-
-        if (acceptDocs != null) {
-          query =
-              new CagraQuery.Builder(getCuVSResourcesInstance())
-                  .withTopK(topK)
-                  .withSearchParams(searchParams)
-                  .withQueryVectors(queryVector)
-                  .withPrefilter(mask[0], maskLength)
-                  .build();
-        } else {
-          query =
+        try (CuVSMatrix queryVector = builder.build()) {
+          CagraQuery.Builder queryBuilder =
               new CagraQuery.Builder(getCuVSResourcesInstance())
                   .withTopK(topK)
                   .withSearchParams(searchParams)
-                  .withQueryVectors(queryVector)
-                  .build();
+                  .withQueryVectors(queryVector);
+          if (acceptDocs != null) {
+            queryBuilder.withPrefilter(mask[0], maskLength);
+          }
+          CagraIndex cagraIndex = cuvsIndex.getCagraIndex();
+          assert cagraIndex != null;
+          searchResult = cagraIndex.search(queryBuilder.build()).getResults();
         }
-        searchResult = cagraIndex.search(query).getResults();
       } else {
         BruteForceIndex bruteforceIndex = cuvsIndex.getBruteforceIndex();
         assert bruteforceIndex != null;