[Backport 3.6] Fix flaky org.opensearch.knn.index.codec.jvector.KNNJVectorTests.testJVectorKnnIndex_mergeEnabled (2nd attempt)

src/main/java/org/opensearch/knn/common/KNNConstants.java

@@ -108,13 +108,6 @@ public class KNNConstants {
     public static final String METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED = "advanced.leading_segment_merge_disabled";
     public static final boolean DEFAULT_LEADING_SEGMENT_MERGE_DISABLED = false;
 
-    // Parameters that only affect search
-    // By default, the entry point is selected more or less "randomly" (due to the concurrent graph construction) which sometimes
-    // causes difficulties getting reproducible, stable results. By disabling entry point selection, the fixed value (0, 0)
-    // will be always selected instead. This setting is not supposed to be used in production.
-    public static final String METHOD_PARAMETER_ENTRY_POINT_SELECTION_DISABLED = "advanced.entry_point_selection_disabled";
-    public static final Boolean DEFAULT_ENTRY_POINT_SELECTION_DISABLED = false;
-
     // API Constants
     public static final String CLEAR_CACHE = "clear_cache";
 

src/main/java/org/opensearch/knn/index/codec/KNN9120Codec/KNN9120PerFieldKnnVectorsFormat.java

@@ -71,8 +71,7 @@ public KNN9120PerFieldKnnVectorsFormat(final Optional<MapperService> mapperServi
                             knnVectorsFormatParams.getNumberOfSubspacesPerVectorSupplier(),
                             knnVectorsFormatParams.getMinBatchSizeForQuantization(),
                             knnVectorsFormatParams.isHierarchyEnabled(),
-                            knnVectorsFormatParams.isLeadingSegmentMergeDisabled(),
-                            knnVectorsFormatParams.isEntryPointSelectionDisabled()
+                            knnVectorsFormatParams.isLeadingSegmentMergeDisabled()
                         );
                     default:
                         throw new IllegalArgumentException("Unsupported java engine: " + knnEngine);

src/main/java/org/opensearch/knn/index/codec/jvector/JVectorFormat.java

@@ -46,7 +46,9 @@ public class JVectorFormat extends KnnVectorsFormat {
     private final float neighborOverflow;
     private final boolean hierarchyEnabled;
     private final boolean leadingSegmentMergeDisabled;
-    private final boolean entryPointSelectionDisabled;
+    private final ForkJoinPool simdPoolMerge;
+    private final ForkJoinPool simdPoolFlush;
+    private final ForkJoinPool parallelismPool;
 
     public JVectorFormat() {
         this(
@@ -59,7 +61,9 @@ public JVectorFormat() {
             KNNConstants.DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION,
             KNNConstants.DEFAULT_HIERARCHY_ENABLED,
             KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED,
-            KNNConstants.DEFAULT_ENTRY_POINT_SELECTION_DISABLED
+            SIMD_POOL_MERGE,
+            SIMD_POOL_FLUSH,
+            PARALLELISM_POOL
         );
     }
 
@@ -68,10 +72,29 @@ public JVectorFormat(int minBatchSizeForQuantization) {
     }
 
     public JVectorFormat(int minBatchSizeForQuantization, boolean leadingSegmentMergeDisabled) {
-        this(minBatchSizeForQuantization, leadingSegmentMergeDisabled, KNNConstants.DEFAULT_ENTRY_POINT_SELECTION_DISABLED);
+        this(
+            NAME,
+            DEFAULT_MAX_CONN,
+            DEFAULT_BEAM_WIDTH,
+            KNNConstants.DEFAULT_NEIGHBOR_OVERFLOW_VALUE.floatValue(),
+            KNNConstants.DEFAULT_ALPHA_VALUE.floatValue(),
+            JVectorFormat::getDefaultNumberOfSubspacesPerVector,
+            minBatchSizeForQuantization,
+            KNNConstants.DEFAULT_HIERARCHY_ENABLED,
+            leadingSegmentMergeDisabled,
+            SIMD_POOL_MERGE,
+            SIMD_POOL_FLUSH,
+            PARALLELISM_POOL
+        );
     }
 
-    public JVectorFormat(int minBatchSizeForQuantization, boolean leadingSegmentMergeDisabled, boolean entryPointSelectionDisabled) {
+    public JVectorFormat(
+        int minBatchSizeForQuantization,
+        boolean leadingSegmentMergeDisabled,
+        final ForkJoinPool simdPoolMerge,
+        final ForkJoinPool simdPoolFlush,
+        final ForkJoinPool parallelismPool
+    ) {
         this(
             NAME,
             DEFAULT_MAX_CONN,
@@ -82,7 +105,9 @@ public JVectorFormat(int minBatchSizeForQuantization, boolean leadingSegmentMerg
             minBatchSizeForQuantization,
             KNNConstants.DEFAULT_HIERARCHY_ENABLED,
             leadingSegmentMergeDisabled,
-            entryPointSelectionDisabled
+            simdPoolMerge,
+            simdPoolFlush,
+            parallelismPool
         );
     }
 
@@ -94,8 +119,7 @@ public JVectorFormat(
         Function<Integer, Integer> numberOfSubspacesPerVectorSupplier,
         int minBatchSizeForQuantization,
         boolean hierarchyEnabled,
-        boolean leadingSegmentMergeDisabled,
-        boolean entryPointSelectionDisabled
+        boolean leadingSegmentMergeDisabled
     ) {
         this(
             NAME,
@@ -107,7 +131,9 @@ public JVectorFormat(
             minBatchSizeForQuantization,
             hierarchyEnabled,
             leadingSegmentMergeDisabled,
-            entryPointSelectionDisabled
+            SIMD_POOL_MERGE,
+            SIMD_POOL_FLUSH,
+            PARALLELISM_POOL
         );
     }
 
@@ -121,7 +147,9 @@ public JVectorFormat(
         int minBatchSizeForQuantization,
         boolean hierarchyEnabled,
         boolean leadingSegmentMergeDisabled,
-        boolean entryPointSelectionDisabled
+        final ForkJoinPool simdPoolMerge,
+        final ForkJoinPool simdPoolFlush,
+        final ForkJoinPool parallelismPool
     ) {
         super(name);
         this.maxConn = maxConn;
@@ -132,7 +160,9 @@ public JVectorFormat(
         this.neighborOverflow = neighborOverflow;
         this.hierarchyEnabled = hierarchyEnabled;
         this.leadingSegmentMergeDisabled = leadingSegmentMergeDisabled;
-        this.entryPointSelectionDisabled = entryPointSelectionDisabled;
+        this.simdPoolMerge = simdPoolMerge;
+        this.simdPoolFlush = simdPoolFlush;
+        this.parallelismPool = parallelismPool;
     }
 
     @Override
@@ -147,7 +177,9 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
             minBatchSizeForQuantization,
             hierarchyEnabled,
             leadingSegmentMergeDisabled,
-            entryPointSelectionDisabled
+            simdPoolMerge,
+            simdPoolFlush,
+            parallelismPool
         );
     }
 

src/main/java/org/opensearch/knn/index/codec/jvector/JVectorWriter.java

@@ -46,9 +46,6 @@
 
 import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED;
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.readVectorEncoding;
-import static org.opensearch.knn.index.codec.jvector.JVectorFormat.PARALLELISM_POOL;
-import static org.opensearch.knn.index.codec.jvector.JVectorFormat.SIMD_POOL_FLUSH;
-import static org.opensearch.knn.index.codec.jvector.JVectorFormat.SIMD_POOL_MERGE;
 
 /**
  * JVectorWriter is responsible for writing vector data into index segments using the JVector library.
@@ -97,8 +94,10 @@ public class JVectorWriter extends KnnVectorsWriter {
     private final int minimumBatchSizeForQuantization; // Threshold for the vector count above which we will trigger PQ quantization
     private final boolean hierarchyEnabled;
     private final boolean leadingSegmentMergeDisabled;
-    // The entry point selection is randomized, allow to disable it for predictable test runs
-    private final boolean entryPointSelectionDisabled;
+
+    private final ForkJoinPool simdPoolMerge;
+    private final ForkJoinPool simdPoolFlush;
+    private final ForkJoinPool parallelismPool;
 
     private boolean finished = false;
 
@@ -112,7 +111,9 @@ public JVectorWriter(
         int minimumBatchSizeForQuantization,
         boolean hierarchyEnabled,
         boolean leadingSegmentMergeDisabled,
-        boolean entryPointSelectionDisabled
+        final ForkJoinPool simdPoolMerge,
+        final ForkJoinPool simdPoolFlush,
+        final ForkJoinPool parallelismPool
     ) throws IOException {
         this.segmentWriteState = segmentWriteState;
         this.maxConn = maxConn;
@@ -123,7 +124,9 @@ public JVectorWriter(
         this.minimumBatchSizeForQuantization = minimumBatchSizeForQuantization;
         this.hierarchyEnabled = hierarchyEnabled;
         this.leadingSegmentMergeDisabled = leadingSegmentMergeDisabled;
-        this.entryPointSelectionDisabled = entryPointSelectionDisabled;
+        this.simdPoolMerge = simdPoolMerge;
+        this.simdPoolFlush = simdPoolFlush;
+        this.parallelismPool = parallelismPool;
 
         String metaFileName = IndexFileNames.segmentFileName(
             segmentWriteState.segmentInfo.name,
@@ -251,7 +254,7 @@ public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
                 randomAccessVectorValues,
                 fieldInfo,
                 segmentWriteState.segmentInfo.name,
-                SIMD_POOL_FLUSH
+                simdPoolFlush
             );
             writeField(field.fieldInfo, randomAccessVectorValues, pqVectors, graphNodeIdToDocMap, graph);
 
@@ -392,7 +395,7 @@ private PQVectors getPQVectors(RandomAccessVectorValues randomAccessVectorValues
             numberOfClustersPerSubspace, // number of centroids per subspace
             vectorSimilarityFunction == VectorSimilarityFunction.EUCLIDEAN, // center the dataset
             UNWEIGHTED,
-            SIMD_POOL_MERGE,
+            simdPoolMerge,
             ForkJoinPool.commonPool()
         );
 
@@ -402,7 +405,7 @@ private PQVectors getPQVectors(RandomAccessVectorValues randomAccessVectorValues
         KNNCounter.KNN_QUANTIZATION_TRAINING_TIME.add(trainingTime);
         log.info("Encoding and building PQ vectors for field {} for {} vectors", fieldName, randomAccessVectorValues.size());
         // PQVectors pqVectors = pq.encodeAll(randomAccessVectorValues, SIMD_POOL);
-        PQVectors pqVectors = PQVectors.encodeAndBuild(pq, randomAccessVectorValues.size(), randomAccessVectorValues, SIMD_POOL_MERGE);
+        PQVectors pqVectors = PQVectors.encodeAndBuild(pq, randomAccessVectorValues.size(), randomAccessVectorValues, simdPoolMerge);
         log.info(
             "Encoded and built PQ vectors for field {}, original size: {} bytes, compressed size: {} bytes",
             fieldName,
@@ -965,7 +968,7 @@ public void merge() throws IOException {
                     compactOrdsToRavvOrds.length,
                     compactOrdsToRavvOrds,
                     this,
-                    SIMD_POOL_MERGE
+                    simdPoolMerge
                 );
             }
 
@@ -982,7 +985,7 @@ public void merge() throws IOException {
                         fieldName
                     );
                     var bsp = BuildScoreProvider.randomAccessScoreProvider(compactRavv, getVectorSimilarityFunction(fieldInfo));
-                    var graph = getGraph(bsp, compactRavv, fieldInfo, segmentWriteState.segmentInfo.name, SIMD_POOL_MERGE);
+                    var graph = getGraph(bsp, compactRavv, fieldInfo, segmentWriteState.segmentInfo.name, simdPoolMerge);
                     writeField(fieldInfo, compactRavv, null, compactOrdToDocMap, graph);
                 }
             } else {
@@ -991,7 +994,7 @@ public void merge() throws IOException {
                 var buildScoreProvider = BuildScoreProvider.pqBuildScoreProvider(getVectorSimilarityFunction(fieldInfo), compactPqVectors);
                 // Pre-init the diversity provider here to avoid doing it lazily (as it could block the SIMD threads)
                 buildScoreProvider.diversityProviderFor(0);
-                var graph = getGraph(buildScoreProvider, compactRavv, fieldInfo, segmentWriteState.segmentInfo.name, SIMD_POOL_MERGE);
+                var graph = getGraph(buildScoreProvider, compactRavv, fieldInfo, segmentWriteState.segmentInfo.name, simdPoolMerge);
                 writeField(fieldInfo, compactRavv, compactPqVectors, compactOrdToDocMap, graph);
             }
         }
@@ -1142,20 +1145,18 @@ private boolean tryLeadingSegmentMerge() throws IOException {
                         degreeOverflow,
                         alpha,
                         true,
-                        SIMD_POOL_MERGE,
-                        PARALLELISM_POOL
+                        simdPoolMerge,
+                        parallelismPool
                     )
                 ) {
                     var vv = heapRavv.threadLocalSupplier();
 
                     // parallel graph construction from the merge documents Ids
-                    SIMD_POOL_MERGE.submit(
-                        () -> IntStream.range(leadingGraph.getIdUpperBound(), heapRavv.size()).parallel().forEach(ord -> {
-                            assert heapToGlobalRavvOrds[ord] != GraphNodeIdToDocMap.NO_VECTOR_OR_DELETED_DOC
-                                : "Should be a valid graph node / vector";
-                            builder.addGraphNode(ord, vv.get().getVector(ord));
-                        })
-                    ).join();
+                    simdPoolMerge.submit(() -> IntStream.range(leadingGraph.getIdUpperBound(), heapRavv.size()).parallel().forEach(ord -> {
+                        assert heapToGlobalRavvOrds[ord] != GraphNodeIdToDocMap.NO_VECTOR_OR_DELETED_DOC
+                            : "Should be a valid graph node / vector";
+                        builder.addGraphNode(ord, vv.get().getVector(ord));
+                    })).join();
 
                     // mark deleted nodes
                     for (int i = 0; i < numBaseVectors; i++) {
@@ -1167,10 +1168,6 @@ private boolean tryLeadingSegmentMerge() throws IOException {
 
                     builder.cleanup();
 
-                    if (entryPointSelectionDisabled == true) {
-                        builder.setEntryPoint(0, 0);
-                    }
-
                     graph = (OnHeapGraphIndex) builder.getGraph();
                 }
 
@@ -1256,10 +1253,6 @@ public OnHeapGraphIndex getGraph(
         })).join();
         graphIndexBuilder.cleanup();
 
-        if (entryPointSelectionDisabled == true) {
-            graphIndexBuilder.setEntryPoint(0, 0);
-        }
-
         graphIndex = (OnHeapGraphIndex) graphIndexBuilder.getGraph();
         final long end = Clock.systemDefaultZone().millis();
 

src/main/java/org/opensearch/knn/index/codec/params/KNNVectorsFormatParams.java

@@ -27,7 +27,6 @@ public class KNNVectorsFormatParams {
     private Function<Integer, Integer> numberOfSubspacesPerVectorSupplier;
     private final SpaceType spaceType;
     private boolean leadingSegmentMergeDisabled;
-    private boolean entryPointSelectionDisabled;
 
     public KNNVectorsFormatParams(final Map<String, Object> params, int defaultMaxConnections, int defaultBeamWidth) {
         this(
@@ -61,7 +60,6 @@ public KNNVectorsFormatParams(
         initNumberOfSubspacesPerVectorSupplier(params);
         this.spaceType = spaceType;
         initLeadingSegmentMergeDisabled(params, KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED);
-        initEntryPointSelectionDisabled(params, KNNConstants.DEFAULT_ENTRY_POINT_SELECTION_DISABLED);
     }
 
     public boolean validate(final Map<String, Object> params) {
@@ -132,12 +130,4 @@ private void initLeadingSegmentMergeDisabled(final Map<String, Object> params, b
         }
         this.leadingSegmentMergeDisabled = defaultLsmDisabled;
     }
-
-    private void initEntryPointSelectionDisabled(final Map<String, Object> params, boolean defaultEpsDisabled) {
-        if (params != null && params.containsKey(KNNConstants.METHOD_PARAMETER_ENTRY_POINT_SELECTION_DISABLED)) {
-            this.entryPointSelectionDisabled = (boolean) params.get(KNNConstants.METHOD_PARAMETER_ENTRY_POINT_SELECTION_DISABLED);
-            return;
-        }
-        this.entryPointSelectionDisabled = defaultEpsDisabled;
-    }
 }