[Java] Refactor SearchResults implementation classes (#1067)

This PR refactors the `SearchResults` implementation classes, by removing inheritance and replacing it with factory methods to create results from native data.
The current implementation classes use inheritance and polymorphism just during creation (in the ctors); this is confusing as it forces to navigate across derived/base class to understand the mechanism, and forces the base class to hold onto references to MemorySegments -- something we don't need to do.

This is preliminary work towards #1037, as it makes the lifetime of those MemorySegment clearer.

Authors:
  - Lorenzo Dematté (https://github.com/ldematte)
  - MithunR (https://github.com/mythrocks)

Approvers:
  - Chris Hegarty (https://github.com/ChrisHegarty)
  - MithunR (https://github.com/mythrocks)

URL: #1067

Author: ldematte

java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java

@@ -15,7 +15,6 @@
  */
 package com.nvidia.cuvs;
 
-import com.nvidia.cuvs.BruteForceIndex.Builder;
 import com.nvidia.cuvs.spi.CuVSProvider;
 import java.io.InputStream;
 import java.io.OutputStream;

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java

@@ -317,7 +317,7 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable {
         checkCuVSError(returnValue, "cuvsRMMFree");
       }
 
-      return new BruteForceSearchResults(
+      return BruteForceSearchResults.create(
           neighborsSequenceLayout,
           distancesSequenceLayout,
           neighborsMemorySegment,

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceSearchResults.java

@@ -15,56 +15,33 @@
  */
 package com.nvidia.cuvs.internal;
 
-import com.nvidia.cuvs.internal.common.SearchResultsImpl;
+import com.nvidia.cuvs.SearchResults;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SequenceLayout;
-import java.util.LinkedHashMap;
-import java.util.Map;
 import java.util.function.LongToIntFunction;
 
 /**
  * SearchResult encapsulates the logic for reading and holding search results.
  *
  * @since 25.02
  */
-public class BruteForceSearchResults extends SearchResultsImpl {
+class BruteForceSearchResults {
 
-  protected BruteForceSearchResults(
+  static SearchResults create(
       SequenceLayout neighboursSequenceLayout,
       SequenceLayout distancesSequenceLayout,
       MemorySegment neighboursMemorySegment,
       MemorySegment distancesMemorySegment,
       int topK,
       LongToIntFunction mapping,
       long numberOfQueries) {
-    super(
+    return SearchResultsImpl.create(
         neighboursSequenceLayout,
         distancesSequenceLayout,
         neighboursMemorySegment,
         distancesMemorySegment,
         topK,
         mapping,
         numberOfQueries);
-    readResultMemorySegments();
-  }
-
-  /**
-   * Reads neighbors and distances {@link MemorySegment} and loads the values
-   * internally
-   */
-  protected void readResultMemorySegments() {
-    Map<Integer, Float> intermediateResultMap = new LinkedHashMap<Integer, Float>();
-    int count = 0;
-    for (long i = 0; i < topK * numberOfQueries; i++) {
-      long id = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i);
-      float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i);
-      intermediateResultMap.put(mapping.applyAsInt(id), dst);
-      count += 1;
-      if (count == topK) {
-        results.add(intermediateResultMap);
-        intermediateResultMap = new LinkedHashMap<Integer, Float>();
-        count = 0;
-      }
-    }
   }
 }

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java

@@ -360,7 +360,7 @@ public SearchResults search(CagraQuery query) throws Throwable {
       returnValue = cuvsRMMFree(cuvsRes, prefilterDP, C_INT_BYTE_SIZE * prefilterBytes);
       checkCuVSError(returnValue, "cuvsRMMFree");
 
-      return new CagraSearchResults(
+      return CagraSearchResults.create(
           neighborsSequenceLayout,
           distancesSequenceLayout,
           neighborsMemorySegment,

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraSearchResults.java

@@ -15,10 +15,13 @@
  */
 package com.nvidia.cuvs.internal;
 
-import com.nvidia.cuvs.internal.common.SearchResultsImpl;
+import com.nvidia.cuvs.SearchResults;
+import java.lang.foreign.MemoryLayout;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SequenceLayout;
 import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.List;
 import java.util.Map;
 import java.util.function.LongToIntFunction;
 
@@ -27,46 +30,44 @@
  *
  * @since 25.02
  */
-public class CagraSearchResults extends SearchResultsImpl {
+class CagraSearchResults {
 
-  protected CagraSearchResults(
+  /**
+   * Factory method to create an on-heap SearchResults (backed by standard Java data types and containers) from
+   * native/off-heap memory data structures.
+   * This class provides its own implementation for reading from native memory instead of reling on
+   * {@link SearchResultsImpl#create} because it requires special handling of neighbours IDs.
+   */
+  static SearchResults create(
       SequenceLayout neighboursSequenceLayout,
       SequenceLayout distancesSequenceLayout,
       MemorySegment neighboursMemorySegment,
       MemorySegment distancesMemorySegment,
       int topK,
       LongToIntFunction mapping,
       long numberOfQueries) {
-    super(
-        neighboursSequenceLayout,
-        distancesSequenceLayout,
-        neighboursMemorySegment,
-        distancesMemorySegment,
-        topK,
-        mapping,
-        numberOfQueries);
-    readResultMemorySegments();
-  }
 
-  /**
-   * Reads neighbors and distances {@link MemorySegment} and loads the values
-   * internally
-   */
-  protected void readResultMemorySegments() {
-    Map<Integer, Float> intermediateResultMap = new LinkedHashMap<Integer, Float>();
+    List<Map<Integer, Float>> results = new LinkedList<>();
+    Map<Integer, Float> intermediateResultMap = new LinkedHashMap<>();
+    var neighboursVarHandle =
+        neighboursSequenceLayout.varHandle(MemoryLayout.PathElement.sequenceElement());
+    var distancesVarHandle =
+        distancesSequenceLayout.varHandle(MemoryLayout.PathElement.sequenceElement());
+
     int count = 0;
     for (long i = 0; i < topK * numberOfQueries; i++) {
-      long id = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i);
+      long id = (long) neighboursVarHandle.get(neighboursMemorySegment, 0, i);
       float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i);
       if (id != Integer.MAX_VALUE) {
         intermediateResultMap.put(mapping.applyAsInt(id), dst);
       }
       count += 1;
       if (count == topK) {
         results.add(intermediateResultMap);
-        intermediateResultMap = new LinkedHashMap<Integer, Float>();
+        intermediateResultMap = new LinkedHashMap<>();
         count = 0;
       }
     }
+    return new SearchResultsImpl(results);
   }
 }

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java

@@ -133,7 +133,7 @@ public SearchResults search(HnswQuery query) throws Throwable {
       returnValue = cuvsStreamSync(cuvsRes);
       checkCuVSError(returnValue, "cuvsStreamSync");
 
-      return new HnswSearchResults(
+      return HnswSearchResults.create(
           neighborsSequenceLayout,
           distancesSequenceLayout,
           neighborsMemorySegment,

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswSearchResults.java

@@ -15,56 +15,33 @@
  */
 package com.nvidia.cuvs.internal;
 
-import com.nvidia.cuvs.internal.common.SearchResultsImpl;
+import com.nvidia.cuvs.SearchResults;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SequenceLayout;
-import java.util.LinkedHashMap;
-import java.util.Map;
 import java.util.function.LongToIntFunction;
 
 /**
  * SearchResult encapsulates the logic for reading and holding search results.
  *
  * @since 25.02
  */
-public class HnswSearchResults extends SearchResultsImpl {
+class HnswSearchResults {
 
-  protected HnswSearchResults(
+  static SearchResults create(
       SequenceLayout neighboursSequenceLayout,
       SequenceLayout distancesSequenceLayout,
       MemorySegment neighboursMemorySegment,
       MemorySegment distancesMemorySegment,
       int topK,
       LongToIntFunction mapping,
       long numberOfQueries) {
-    super(
+    return SearchResultsImpl.create(
         neighboursSequenceLayout,
         distancesSequenceLayout,
         neighboursMemorySegment,
         distancesMemorySegment,
         topK,
         mapping,
         numberOfQueries);
-    readResultMemorySegments();
-  }
-
-  /**
-   * Reads neighbors and distances {@link MemorySegment} and loads the values
-   * internally
-   */
-  protected void readResultMemorySegments() {
-    Map<Integer, Float> intermediateResultMap = new LinkedHashMap<Integer, Float>();
-    int count = 0;
-    for (long i = 0; i < topK * numberOfQueries; i++) {
-      long id = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i);
-      float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i);
-      intermediateResultMap.put(mapping.applyAsInt(id), dst);
-      count += 1;
-      if (count == topK) {
-        results.add(intermediateResultMap);
-        intermediateResultMap = new LinkedHashMap<Integer, Float>();
-        count = 0;
-      }
-    }
   }
 }

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/SearchResultsImpl.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.nvidia.cuvs.internal;
+
+import com.nvidia.cuvs.SearchResults;
+import java.lang.foreign.MemoryLayout;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SequenceLayout;
+import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.function.LongToIntFunction;
+
+class SearchResultsImpl implements SearchResults {
+
+  private final List<Map<Integer, Float>> results;
+
+  SearchResultsImpl(List<Map<Integer, Float>> results) {
+    this.results = results;
+  }
+
+  /**
+   * Factory method to create an on-heap SearchResults (backed by standard Java data types and containers) from
+   * native/off-heap memory data structures.
+   */
+  static SearchResults create(
+      SequenceLayout neighboursSequenceLayout,
+      SequenceLayout distancesSequenceLayout,
+      MemorySegment neighboursMemorySegment,
+      MemorySegment distancesMemorySegment,
+      int topK,
+      LongToIntFunction mapping,
+      long numberOfQueries) {
+    List<Map<Integer, Float>> results = new LinkedList<>();
+    Map<Integer, Float> intermediateResultMap = new LinkedHashMap<>();
+    var neighboursVarHandle =
+        neighboursSequenceLayout.varHandle(MemoryLayout.PathElement.sequenceElement());
+    var distancesVarHandle =
+        distancesSequenceLayout.varHandle(MemoryLayout.PathElement.sequenceElement());
+
+    int count = 0;
+    for (long i = 0; i < topK * numberOfQueries; i++) {
+      long id = (long) neighboursVarHandle.get(neighboursMemorySegment, 0L, i);
+      float dst = (float) distancesVarHandle.get(distancesMemorySegment, 0L, i);
+      intermediateResultMap.put(mapping != null ? mapping.applyAsInt((int) id) : (int) id, dst);
+      count += 1;
+      if (count == topK) {
+        results.add(intermediateResultMap);
+        intermediateResultMap = new LinkedHashMap<>();
+        count = 0;
+      }
+    }
+
+    return new SearchResultsImpl(results);
+  }
+
+  /**
+   * Gets a list results as a map of neighbor IDs to distances.
+   *
+   * @return a list of results for each query as a map of neighbor IDs to distance
+   */
+  @Override
+  public List<Map<Integer, Float>> getResults() {
+    return results;
+  }
+}