Fix handling of deletes during leading segment merge

CHANGELOG.md

@@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Bug Fixes
 * Fix visited docs tracking that led to assertions on AbstractKnnVectorQuery side [238] (https://github.com/opensearch-project/opensearch-jvector/pull/238)
 * Revert of support deletes for incremental insertion [240](https://github.com/opensearch-project/opensearch-jvector/pull/240)
+* Fix leading segment merge for deletions [242](https://github.com/opensearch-project/opensearch-jvector/pull/242)
 ### Infrastructure
 * Upgrade Gradle to 9.2.0 [220] (https://github.com/opensearch-project/opensearch-jvector/pull/222)
 * Add support for JDK25 [220] (https://github.com/opensearch-project/opensearch-jvector/pull/222)

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/DelayedInitDiversityProvider.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.codec.jvector;
+
+import org.opensearch.common.SetOnce;
+import org.opensearch.common.SetOnce.AlreadySetException;
+
+import io.github.jbellis.jvector.graph.NodeArray;
+import io.github.jbellis.jvector.graph.diversity.DiversityProvider;
+import io.github.jbellis.jvector.util.BitSet;
+import lombok.NonNull;
+
+/**
+ * <p>A diversity provider that can be initialized after construction.
+ *
+ * <p>
+ * This is useful during leading segment merge when loading the
+ * neighbors score cache from disk. The score cache is loaded as an
+ * OnHeapGraphIndex, which requires a diveristy provider during construction.
+ * However, creating the diversity provider requires knowledge of the ordinals
+ * contained in the graph which is only available after the graph is loaded.
+ * Since the diversity provider is not used until the graph is mutated, this
+ * class can be used to construct the graph without a diversity provider
+ * available.
+ */
+class DelayedInitDiversityProvider implements DiversityProvider {
+
+    private SetOnce<DiversityProvider> delegate = new SetOnce<>();
+
+    /**
+     * Creates an uninitialized diversity provider.
+     * Call {@link #initialize} before use.
+     */
+    DelayedInitDiversityProvider() {}
+
+    /**
+     * Initialize this DiversityProvider with a delegate
+     * @throws AlreadySetException if already set
+     */
+    void initialize(@NonNull DiversityProvider delegate) throws AlreadySetException {
+        this.delegate.set(delegate);
+    }
+
+    @Override
+    public double retainDiverse(NodeArray neighbors, int maxDegree, int diverseBefore, BitSet selected) {
+        if (delegate.get() == null) {
+            throw new IllegalStateException("DelayedInitDiversityProvider was not initialzied (call initialize() before use)");
+        }
+        return delegate.get().retainDiverse(neighbors, maxDegree, diverseBefore, selected);
+    }
+}

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

@@ -43,10 +43,15 @@ public GraphNodeIdToDocMap(IndexInput in) throws IOException {
 
         graphNodeIdsToDocIds = new int[size];
         docIdsToGraphNodeIds = new int[maxDocId];
+        Arrays.fill(graphNodeIdsToDocIds, -1);
+        Arrays.fill(docIdsToGraphNodeIds, -1);
         for (int ord = 0; ord < size; ord++) {
             final int docId = in.readVInt();
             graphNodeIdsToDocIds[ord] = docId;
-            docIdsToGraphNodeIds[docId] = ord;
+            if (docId != -1) {
+                // ignore deleted documents
+                docIdsToGraphNodeIds[docId] = ord;
+            }
         }
     }
 
@@ -67,19 +72,21 @@ 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.
-        if (maxDocs < graphNodeIdsToDocIds.length) {
-            throw new IllegalStateException("Max docs " + maxDocs + " is less than the number of ordinals " + graphNodeIdsToDocIds.length);
-        }
-        if (maxDocId > graphNodeIdsToDocIds.length) {
-            log.warn(
-                "Max doc id {} is greater than the number of ordinals {}, this implies a lot of deleted documents. Or that some documents are missing vectors. Wasting a lot of memory",
-                maxDocId,
+        // 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 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
             );
         }
         this.docIdsToGraphNodeIds = new int[maxDocs];
         Arrays.fill(this.docIdsToGraphNodeIds, -1); // -1 means no mapping to ordinal
         for (int ord = 0; ord < graphNodeIdsToDocIds.length; ord++) {
+            // -1 means no mapping to docId since document was deleted
+            if (graphNodeIdsToDocIds[ord] == -1) {
+                continue;
+            }
             this.docIdsToGraphNodeIds[graphNodeIdsToDocIds[ord]] = ord;
         }
     }

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

@@ -36,6 +36,7 @@ public class JVectorFormat extends KnnVectorsFormat {
     // Unfortunately, this can't be managed yet by the OpenSearch ThreadPool because it's not supporting {@link ForkJoinPool} types
     public static final ForkJoinPool SIMD_POOL_MERGE = getPhysicalCoreExecutor();
     public static final ForkJoinPool SIMD_POOL_FLUSH = getPhysicalCoreExecutor();
+    public static final ForkJoinPool PARALLELISM_POOL = ForkJoinPool.commonPool();
 
     private final int maxConn;
     private final int beamWidth;
@@ -44,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(
@@ -54,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,
@@ -67,7 +74,8 @@ public JVectorFormat(int minBatchSizeForQuantization) {
             KNNConstants.DEFAULT_ALPHA_VALUE.floatValue(),
             JVectorFormat::getDefaultNumberOfSubspacesPerVector,
             minBatchSizeForQuantization,
-            KNNConstants.DEFAULT_HIERARCHY_ENABLED
+            KNNConstants.DEFAULT_HIERARCHY_ENABLED,
+            leadingSegmentMergeDisabled
         );
     }
 
@@ -78,7 +86,8 @@ public JVectorFormat(
         float alpha,
         Function<Integer, Integer> numberOfSubspacesPerVectorSupplier,
         int minBatchSizeForQuantization,
-        boolean hierarchyEnabled
+        boolean hierarchyEnabled,
+        boolean leadingSegmentMergeDisabled
     ) {
         this(
             NAME,
@@ -88,7 +97,8 @@ public JVectorFormat(
             alpha,
             numberOfSubspacesPerVectorSupplier,
             minBatchSizeForQuantization,
-            hierarchyEnabled
+            hierarchyEnabled,
+            leadingSegmentMergeDisabled
         );
     }
 
@@ -100,7 +110,8 @@ public JVectorFormat(
         float alpha,
         Function<Integer, Integer> numberOfSubspacesPerVectorSupplier,
         int minBatchSizeForQuantization,
-        boolean hierarchyEnabled
+        boolean hierarchyEnabled,
+        boolean leadingSegmentMergeDisabled
     ) {
         super(name);
         this.maxConn = maxConn;
@@ -110,6 +121,7 @@ public JVectorFormat(
         this.alpha = alpha;
         this.neighborOverflow = neighborOverflow;
         this.hierarchyEnabled = hierarchyEnabled;
+        this.leadingSegmentMergeDisabled = leadingSegmentMergeDisabled;
     }
 
     @Override
@@ -122,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/JVectorReader.java

@@ -173,7 +173,8 @@ public void search(String field, float[] target, KnnCollector knnCollector, Acce
             if (acceptDocs == null) compatibleBits = ord -> true;
             else {
                 Bits b = acceptDocs.bits();
-                compatibleBits = ord -> b == null || b.get(jvectorLuceneDocMap.getLuceneDocId(ord));
+                compatibleBits = ord -> b == null
+                    || (jvectorLuceneDocMap.getLuceneDocId(ord) != -1 && b.get(jvectorLuceneDocMap.getLuceneDocId(ord)));
             }
 
             try (var graphSearcher = new GraphSearcher(index)) {