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

@@ -0,0 +1,53 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.codec.jvector;
+
+import java.util.Optional;
+
+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 Optional<DiversityProvider> delegate = Optional.empty();
+
+    /**
+     * Creates an uninitialized diversity provider.
+     * Call {@link #initialize} before use.
+     */
+    DelayedInitDiversityProvider() {}
+
+    /** Initialize this DiversityProvider with a delegate */
+    void initialize(@NonNull DiversityProvider delegate) {
+        if (this.delegate.isPresent()) {
+            throw new IllegalStateException("already initialized");
+        }
+        this.delegate = Optional.of(delegate);
+    }
+
+    @Override
+    public double retainDiverse(NodeArray neighbors, int maxDegree, int diverseBefore, BitSet selected) {
+        if (delegate.isEmpty()) {
+            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,20 @@ 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,
+        if (maxDocs < 0.8 * graphNodeIdsToDocIds.length) {
+            log.info(
+                "Max docs {} is less than the number of ordinals {}, this implies a lot of deleted documents. Or that some documents are missing vectors. Wasting a lot 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/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;

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)) {