Add config to disable leading segment merge

Also tweak some logs

Signed-off-by: Ashwin Krishna Kumar <nebulousmagneticwind@outlook.com>

Author: ashkrisk

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

@@ -100,6 +100,10 @@ public class KNNConstants {
                                                                                 // default
     public static final Boolean DEFAULT_HIERARCHY_ENABLED = false;
 
+    // Parameters that only affect merge
+    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;
+
     // API Constants
     public static final String CLEAR_CACHE = "clear_cache";
 

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

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

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

@@ -72,9 +72,10 @@ public GraphNodeIdToDocMap(int[] graphNodeIdsToDocIds) {
         final int maxDocs = maxDocId + 1;
         // We are going to assume that the number of ordinals is roughly the same as the number of documents in the segment, therefore,
         // the mapping will not be sparse.
+        // Note that the merge process may create ephemeral mappings that are sparse.
         if (maxDocs < 0.8 * graphNodeIdsToDocIds.length) {
             log.info(
-                "Max docs {} is less than 80% the number of ordinals {}, this implies a lot of deleted documents. Or that some documents are missing vectors. Wasting a lot of memory",
+                "Max docs {} is less than 80% the number of ordinals {}. This is normal if many docs were recently deleted or overwritten. Otherwise if many docs are missing vectors this is a waste of memory.",
                 maxDocs,
                 graphNodeIdsToDocIds.length
             );

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

@@ -85,6 +85,14 @@ private static MethodComponent initMethodComponent() {
                     (v, context) -> v != null && v > 0 && v <= context.getDimension()
                 )
             )
+            .addParameter(
+                METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED,
+                new Parameter.BooleanParameter(
+                    METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED,
+                    DEFAULT_LEADING_SEGMENT_MERGE_DISABLED,
+                    (v, context) -> true
+                )
+            )
             .build();
     }
 

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

@@ -45,6 +45,7 @@ public class JVectorFormat extends KnnVectorsFormat {
     private final float alpha;
     private final float neighborOverflow;
     private final boolean hierarchyEnabled;
+    private final boolean leadingSegmentMergeDisabled;
 
     public JVectorFormat() {
         this(
@@ -55,11 +56,16 @@ public JVectorFormat() {
             KNNConstants.DEFAULT_ALPHA_VALUE.floatValue(),
             JVectorFormat::getDefaultNumberOfSubspacesPerVector,
             KNNConstants.DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION,
-            KNNConstants.DEFAULT_HIERARCHY_ENABLED
+            KNNConstants.DEFAULT_HIERARCHY_ENABLED,
+            KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED
         );
     }
 
     public JVectorFormat(int minBatchSizeForQuantization) {
+        this(minBatchSizeForQuantization, KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED);
+    }
+
+    public JVectorFormat(int minBatchSizeForQuantization, boolean leadingSegmentMergeDisabled) {
         this(
             NAME,
             DEFAULT_MAX_CONN,
@@ -68,7 +74,8 @@ public JVectorFormat(int minBatchSizeForQuantization) {
             KNNConstants.DEFAULT_ALPHA_VALUE.floatValue(),
             JVectorFormat::getDefaultNumberOfSubspacesPerVector,
             minBatchSizeForQuantization,
-            KNNConstants.DEFAULT_HIERARCHY_ENABLED
+            KNNConstants.DEFAULT_HIERARCHY_ENABLED,
+            leadingSegmentMergeDisabled
         );
     }
 
@@ -79,7 +86,8 @@ public JVectorFormat(
         float alpha,
         Function<Integer, Integer> numberOfSubspacesPerVectorSupplier,
         int minBatchSizeForQuantization,
-        boolean hierarchyEnabled
+        boolean hierarchyEnabled,
+        boolean leadingSegmentMergeDisabled
     ) {
         this(
             NAME,
@@ -89,7 +97,8 @@ public JVectorFormat(
             alpha,
             numberOfSubspacesPerVectorSupplier,
             minBatchSizeForQuantization,
-            hierarchyEnabled
+            hierarchyEnabled,
+            leadingSegmentMergeDisabled
         );
     }
 
@@ -101,7 +110,8 @@ public JVectorFormat(
         float alpha,
         Function<Integer, Integer> numberOfSubspacesPerVectorSupplier,
         int minBatchSizeForQuantization,
-        boolean hierarchyEnabled
+        boolean hierarchyEnabled,
+        boolean leadingSegmentMergeDisabled
     ) {
         super(name);
         this.maxConn = maxConn;
@@ -111,6 +121,7 @@ public JVectorFormat(
         this.alpha = alpha;
         this.neighborOverflow = neighborOverflow;
         this.hierarchyEnabled = hierarchyEnabled;
+        this.leadingSegmentMergeDisabled = leadingSegmentMergeDisabled;
     }
 
     @Override
@@ -123,7 +134,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
             alpha,
             numberOfSubspacesPerVectorSupplier,
             minBatchSizeForQuantization,
-            hierarchyEnabled
+            hierarchyEnabled,
+            leadingSegmentMergeDisabled
         );
     }
 

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

@@ -96,6 +96,7 @@ public class JVectorWriter extends KnnVectorsWriter {
                                                                                  // as a function of the original dimension
     private final int minimumBatchSizeForQuantization; // Threshold for the vector count above which we will trigger PQ quantization
     private final boolean hierarchyEnabled;
+    private final boolean leadingSegmentMergeDisabled;
 
     private boolean finished = false;
 
@@ -107,7 +108,8 @@ public JVectorWriter(
         float alpha,
         Function<Integer, Integer> numberOfSubspacesPerVectorSupplier,
         int minimumBatchSizeForQuantization,
-        boolean hierarchyEnabled
+        boolean hierarchyEnabled,
+        boolean leadingSegmentMergeDisabled
     ) throws IOException {
         this.segmentWriteState = segmentWriteState;
         this.maxConn = maxConn;
@@ -117,6 +119,8 @@ public JVectorWriter(
         this.numberOfSubspacesPerVectorSupplier = numberOfSubspacesPerVectorSupplier;
         this.minimumBatchSizeForQuantization = minimumBatchSizeForQuantization;
         this.hierarchyEnabled = hierarchyEnabled;
+        this.leadingSegmentMergeDisabled = leadingSegmentMergeDisabled;
+
         String metaFileName = IndexFileNames.segmentFileName(
             segmentWriteState.segmentInfo.name,
             segmentWriteState.segmentSuffix,
@@ -796,7 +800,7 @@ public RandomAccessMergedFloatVectorValues(FieldInfo fieldInfo, MergeState merge
 
                 for (int docId = it.nextDoc(); docId != DocIdSetIterator.NO_MORE_DOCS; docId = it.nextDoc()) {
                     if (docMaps[readerIdx].get(docId) == -1) {
-                        log.warn(
+                        log.debug(
                             "Document {} in reader {} is not mapped to a global ordinal from the merge docMaps. Will skip this document for now",
                             docId,
                             readerIdx
@@ -956,11 +960,21 @@ public void merge() throws IOException {
         /**
          * <p>Perform leading segment merge.
          *
-         * <p>In some cases leading segment merge should be skipped, and this method will return false.
+         * <p>
+         * In some cases leading segment merge should be skipped, and this method will return false.
+         * This is the case when:
+         * - Leading segment merge is disabled through configuration
+         * - There is a risk of integer overflow due to sparsity of the OnHeapGraph
+         * - The OnHeapGraph is too sparse relative to it's size
          *
          * @return a boolean value indicating if leading segment merge was performed
          */
         private boolean tryLeadingSegmentMerge() throws IOException {
+            if (leadingSegmentMergeDisabled) {
+                log.info("Leading segment merge is disabled, skipping");
+                return false;
+            }
+
             var leadingFieldsReader = (PerFieldKnnVectorsFormat.FieldsReader) readers[LEADING_READER_IDX];
             var leadingReader = (JVectorReader) leadingFieldsReader.getFieldReader(fieldInfo.name);
             var graphReader = leadingReader.getNeighborsScoreCacheForField(fieldInfo.name);

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

@@ -26,6 +26,7 @@ public class KNNVectorsFormatParams {
     private boolean hierarchyEnabled;
     private Function<Integer, Integer> numberOfSubspacesPerVectorSupplier;
     private final SpaceType spaceType;
+    private boolean leadingSegmentMergeDisabled;
 
     public KNNVectorsFormatParams(final Map<String, Object> params, int defaultMaxConnections, int defaultBeamWidth) {
         this(
@@ -58,6 +59,7 @@ public KNNVectorsFormatParams(
         initHierarchyEnabled(params, defaultHierarchyEnabled);
         initNumberOfSubspacesPerVectorSupplier(params);
         this.spaceType = spaceType;
+        initLeadingSegmentMergeDisabled(params, KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED);
     }
 
     public boolean validate(final Map<String, Object> params) {
@@ -120,4 +122,12 @@ private void initNumberOfSubspacesPerVectorSupplier(final Map<String, Object> pa
         }
         this.numberOfSubspacesPerVectorSupplier = JVectorFormat::getDefaultNumberOfSubspacesPerVector;
     }
+
+    private void initLeadingSegmentMergeDisabled(final Map<String, Object> params, boolean defaultLsmDisabled) {
+        if (params != null && params.containsKey(KNNConstants.METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED)) {
+            this.hierarchyEnabled = (boolean) params.get(KNNConstants.METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED);
+            return;
+        }
+        this.leadingSegmentMergeDisabled = defaultLsmDisabled;
+    }
 }

src/test/java/org/opensearch/knn/index/codec/jvector/JVectorWriterMergeTests.java

@@ -85,6 +85,8 @@ static class MergeTestScenario {
         int overqueryFactor = KNNConstants.DEFAULT_OVER_QUERY_FACTOR;
         @Default
         double minimumRecall = 0.99;
+        @Default
+        boolean leadingSegmentMergeDisabled = KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED;
     }
 
     @Rule
@@ -102,7 +104,7 @@ void runScenario(MergeTestScenario scenario) throws IOException {
         // Path indexPath = createTempDir();
         IndexWriterConfig iwc = LuceneTestCase.newIndexWriterConfig();
         iwc.setUseCompoundFile(false);
-        iwc.setCodec(getCodec(scenario.minPqThreshold));
+        iwc.setCodec(getCodec(scenario.minPqThreshold, scenario.leadingSegmentMergeDisabled));
         iwc.setMergePolicy(new ForceMergesOnlyMergePolicy(false));
 
         try (var fsd = FSDirectory.open(tempDir.getRoot().toPath()); var writer = new IndexWriter(fsd, iwc);) {
@@ -350,8 +352,9 @@ public void multiPhaseMergeWithDeletesProgressivePQ() throws IOException {
                     .deletionRanges(List.of(new DeletionRange(998, 1023), new DeletionRange(1023, 1100), new DeletionRange(2502, 2504)))
                     .build()
             )
+            .round(MergeTestRound.builder().segmentSizes(List.of(400)).build())  // an extra round just because
             .overqueryFactor(20)
-            .minimumRecall(0.99)
+            .minimumRecall(0.98)
             .build();
         runScenario(scenario);
     }
@@ -390,4 +393,50 @@ public void multiPhaseMergeWithDeletesAlwaysPQ() throws IOException {
             .build();
         runScenario(scenario);
     }
+
+    @Test
+    public void testLeadingSegmentMergeDisabled() throws IOException {
+        // we'll progressively increase to beyond the PQ threshold
+        var scenario = MergeTestScenario.builder()
+            .minPqThreshold(1000)   // start using PQ once the vector count crosses this
+            .leadingSegmentMergeDisabled(true)  // and disable leading segment merge
+            .round(
+                MergeTestRound.builder()
+                    .segmentSizes(List.of(10, 200, 50, 100, 100))  // count 460
+                    .deletionRanges(
+                        List.of(
+                            // total deletions 126
+                            new DeletionRange(0, 5),      // count 5
+                            new DeletionRange(210, 260),  // count 50
+                            new DeletionRange(330, 400),  // count 70
+                            new DeletionRange(410, 411)   // count 1
+                        )
+                    )
+                    .build()  // count 334
+            )
+            .round(
+                MergeTestRound.builder()
+                    .segmentSizes(List.of(10, 20, 300))  // count 330
+                    .deletionRanges(
+                        List.of(
+                            // total deletions 195
+                            new DeletionRange(405, 600),  // count 195, BUT this range has some overlap so actually 194
+                            new DeletionRange(10, 11)     // count 1
+                        )
+                    )
+                    .build()  // +135, total count 469
+            )
+            .round(MergeTestRound.builder().segmentSizes(List.of(50, 200)).build())  // bonus round
+            .round(
+                MergeTestRound.builder()
+                    .segmentSizes(List.of(20, 2000, 1))  // Add 2000 vectors to ensure PQ
+                    .deletionRanges(List.of(new DeletionRange(998, 1023), new DeletionRange(1023, 1100), new DeletionRange(2502, 2504)))
+                    .build()
+            )
+            .round(MergeTestRound.builder().segmentSizes(List.of(400)).build())
+            .overqueryFactor(20)
+            .minimumRecall(0.99)
+            .build();
+        runScenario(scenario);
+    }
 }

src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java

@@ -49,6 +49,7 @@
 import static org.opensearch.knn.common.KNNConstants.DISK_ANN;
 import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD;
 import static org.opensearch.knn.index.KNNSettings.KNN_INDEX;
+import static org.opensearch.knn.common.KNNConstants.DEFAULT_LEADING_SEGMENT_MERGE_DISABLED;
 import static org.opensearch.knn.common.KNNConstants.DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION;
 
 public class CommonTestUtils {
@@ -123,18 +124,22 @@ public static String createIndexMapping(int dimension, SpaceType spaceType, Vect
     }
 
     public static Codec getCodec() {
-        return getCodec(DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION);
+        return getCodec(DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION, DEFAULT_LEADING_SEGMENT_MERGE_DISABLED);
     }
 
     public static Codec getCodec(int minBatchSizeForQuantization) {
+        return getCodec(minBatchSizeForQuantization, DEFAULT_LEADING_SEGMENT_MERGE_DISABLED);
+    }
+
+    public static Codec getCodec(int minBatchSizeForQuantization, boolean leadingSegmentMergeDisabled) {
         return new FilterCodec(KNNCodecVersion.V_10_03_0.getCodecName(), new Lucene103Codec()) {
             @Override
             public KnnVectorsFormat knnVectorsFormat() {
                 return new PerFieldKnnVectorsFormat() {
 
                     @Override
                     public KnnVectorsFormat getKnnVectorsFormatForField(String field) {
-                        return new JVectorFormat(minBatchSizeForQuantization);
+                        return new JVectorFormat(minBatchSizeForQuantization, leadingSegmentMergeDisabled);
                     }
                 };
             }

src/test/java/org/opensearch/knn/index/engine/JVectorEngineIT.java

@@ -104,6 +104,7 @@ public void testAddDoc() throws Exception {
             .field(KNNConstants.METHOD_PARAMETER_NEIGHBOR_OVERFLOW, 2.0)
             .field(KNNConstants.METHOD_PARAMETER_HIERARCHY_ENABLED, false)
             .field(KNNConstants.METHOD_PARAMETER_MIN_BATCH_SIZE_FOR_QUANTIZATION, 1000)
+            .field(KNNConstants.METHOD_PARAMETER_LEADING_SEGMENT_MERGE_DISABLED, false)
             .endObject()
             .endObject()
             .endObject()