LUCENE-10375: Write merged vectors to file before building graph (#601)

When merging segments together, the `KnnVectorsWriter` creates a `VectorValues`
instance with a merged view of all the segments' vectors. This merged instance
is used when constructing the new HNSW graph. Graph building needs random
access, and the merged VectorValues support this by mapping from merged
ordinals to segments and segment ordinals. This mapping can add significant
overhead when building the graph.

This change updates the HNSW merging logic to first write the combined segment
vectors to a file, then use that the file to build the graph. This helps speed
up segment merging, and also lets us simplify `VectorValuesMerger`, which
provides the merged view of vector values.

Author: jtibshirani

lucene/CHANGES.txt

@@ -128,6 +128,9 @@ Optimizations
 * LUCENE-10379: Count directly into the dense values array in FastTaxonomyFacetCounts#countAll.
   (Guo Feng, Greg Miller)
 
+* LUCENE-10375: Speed up HNSW vectors merge by first writing combined vector
+  data to a file. (Julie Tibshirani, Adrien Grand)
+
 Changes in runtime behavior
 ---------------------
 

lucene/core/src/java/org/apache/lucene/codecs/KnnVectorsWriter.java

@@ -17,19 +17,13 @@
 
 package org.apache.lucene.codecs;
 
-import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
-
 import java.io.Closeable;
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 import org.apache.lucene.index.DocIDMerger;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.MergeState;
-import org.apache.lucene.index.RandomAccessVectorValues;
-import org.apache.lucene.index.RandomAccessVectorValuesProducer;
-import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.index.VectorValues;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.util.Bits;
@@ -48,7 +42,11 @@ public abstract void writeField(FieldInfo fieldInfo, KnnVectorsReader knnVectors
   /** Called once at the end before close */
   public abstract void finish() throws IOException;
 
-  /** Merge the vector values from multiple segments, for all fields */
+  /**
+   * Merges the segment vectors for all fields. This default implementation delegates to {@link
+   * #writeField}, passing a {@link KnnVectorsReader} that combines the vector values and ignores
+   * deleted documents.
+   */
   public void merge(MergeState mergeState) throws IOException {
     for (int i = 0; i < mergeState.fieldInfos.length; i++) {
       KnnVectorsReader reader = mergeState.knnVectorsReaders[i];
@@ -57,163 +55,106 @@ public void merge(MergeState mergeState) throws IOException {
         reader.checkIntegrity();
       }
     }
+
     for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) {
       if (fieldInfo.hasVectorValues()) {
-        mergeVectors(fieldInfo, mergeState);
-      }
-    }
-    finish();
-  }
+        if (mergeState.infoStream.isEnabled("VV")) {
+          mergeState.infoStream.message("VV", "merging " + mergeState.segmentInfo);
+        }
 
-  private void mergeVectors(FieldInfo mergeFieldInfo, final MergeState mergeState)
-      throws IOException {
-    if (mergeState.infoStream.isEnabled("VV")) {
-      mergeState.infoStream.message("VV", "merging " + mergeState.segmentInfo);
-    }
-    // Create a new VectorValues by iterating over the sub vectors, mapping the resulting
-    // docids using docMaps in the mergeState.
-    writeField(
-        mergeFieldInfo,
-        new KnnVectorsReader() {
-          @Override
-          public long ramBytesUsed() {
-            return 0;
-          }
+        writeField(
+            fieldInfo,
+            new KnnVectorsReader() {
+              @Override
+              public long ramBytesUsed() {
+                return 0;
+              }
 
-          @Override
-          public void close() throws IOException {
-            throw new UnsupportedOperationException();
-          }
+              @Override
+              public void close() {
+                throw new UnsupportedOperationException();
+              }
 
-          @Override
-          public void checkIntegrity() throws IOException {
-            throw new UnsupportedOperationException();
-          }
+              @Override
+              public void checkIntegrity() {
+                throw new UnsupportedOperationException();
+              }
 
-          @Override
-          public VectorValues getVectorValues(String field) throws IOException {
-            List<VectorValuesSub> subs = new ArrayList<>();
-            int dimension = -1;
-            VectorSimilarityFunction similarityFunction = null;
-            int nonEmptySegmentIndex = 0;
-            for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) {
-              KnnVectorsReader knnVectorsReader = mergeState.knnVectorsReaders[i];
-              if (knnVectorsReader != null) {
-                if (mergeFieldInfo != null && mergeFieldInfo.hasVectorValues()) {
-                  int segmentDimension = mergeFieldInfo.getVectorDimension();
-                  VectorSimilarityFunction segmentSimilarityFunction =
-                      mergeFieldInfo.getVectorSimilarityFunction();
-                  if (dimension == -1) {
-                    dimension = segmentDimension;
-                    similarityFunction = mergeFieldInfo.getVectorSimilarityFunction();
-                  } else if (dimension != segmentDimension) {
-                    throw new IllegalStateException(
-                        "Varying dimensions for vector-valued field "
-                            + mergeFieldInfo.name
-                            + ": "
-                            + dimension
-                            + "!="
-                            + segmentDimension);
-                  } else if (similarityFunction != segmentSimilarityFunction) {
-                    throw new IllegalStateException(
-                        "Varying similarity functions for vector-valued field "
-                            + mergeFieldInfo.name
-                            + ": "
-                            + similarityFunction
-                            + "!="
-                            + segmentSimilarityFunction);
-                  }
-                  VectorValues values = knnVectorsReader.getVectorValues(mergeFieldInfo.name);
-                  if (values != null) {
-                    subs.add(
-                        new VectorValuesSub(nonEmptySegmentIndex++, mergeState.docMaps[i], values));
-                  }
-                }
+              @Override
+              public VectorValues getVectorValues(String field) throws IOException {
+                return MergedVectorValues.mergeVectorValues(fieldInfo, mergeState);
               }
-            }
-            return new VectorValuesMerger(subs, mergeState);
-          }
 
-          @Override
-          public TopDocs search(String field, float[] target, int k, Bits acceptDocs)
-              throws IOException {
-            throw new UnsupportedOperationException();
-          }
-        });
+              @Override
+              public TopDocs search(String field, float[] target, int k, Bits acceptDocs) {
+                throw new UnsupportedOperationException();
+              }
+            });
 
-    if (mergeState.infoStream.isEnabled("VV")) {
-      mergeState.infoStream.message("VV", "merge done " + mergeState.segmentInfo);
+        if (mergeState.infoStream.isEnabled("VV")) {
+          mergeState.infoStream.message("VV", "merge done " + mergeState.segmentInfo);
+        }
+      }
     }
+    finish();
   }
 
   /** Tracks state of one sub-reader that we are merging */
   private static class VectorValuesSub extends DocIDMerger.Sub {
 
     final VectorValues values;
-    final int segmentIndex;
-    int count;
 
-    VectorValuesSub(int segmentIndex, MergeState.DocMap docMap, VectorValues values) {
+    VectorValuesSub(MergeState.DocMap docMap, VectorValues values) {
       super(docMap);
       this.values = values;
-      this.segmentIndex = segmentIndex;
       assert values.docID() == -1;
     }
 
     @Override
     public int nextDoc() throws IOException {
-      int docId = values.nextDoc();
-      if (docId != NO_MORE_DOCS) {
-        // Note: this does count deleted docs since they are present in the to-be-merged segment
-        ++count;
-      }
-      return docId;
+      return values.nextDoc();
     }
   }
 
-  /**
-   * View over multiple VectorValues supporting iterator-style access via DocIdMerger. Maintains a
-   * reverse ordinal mapping for documents having values in order to support random access by dense
-   * ordinal.
-   */
-  private static class VectorValuesMerger extends VectorValues
-      implements RandomAccessVectorValuesProducer {
+  /** View over multiple VectorValues supporting iterator-style access via DocIdMerger. */
+  public static class MergedVectorValues extends VectorValues {
     private final List<VectorValuesSub> subs;
     private final DocIDMerger<VectorValuesSub> docIdMerger;
-    private final int[] ordBase;
     private final int cost;
-    private int size;
+    private final int size;
 
     private int docId;
     private VectorValuesSub current;
-    /* For each doc with a vector, record its ord in the segments being merged. This enables random
-     * access into the unmerged segments using the ords from the merged segment.
-     */
-    private int[] ordMap;
-    private int ord;
 
-    VectorValuesMerger(List<VectorValuesSub> subs, MergeState mergeState) throws IOException {
+    /** Returns a merged view over all the segment's {@link VectorValues}. */
+    public static MergedVectorValues mergeVectorValues(FieldInfo fieldInfo, MergeState mergeState)
+        throws IOException {
+      assert fieldInfo != null && fieldInfo.hasVectorValues();
+
+      List<VectorValuesSub> subs = new ArrayList<>();
+      for (int i = 0; i < mergeState.knnVectorsReaders.length; i++) {
+        KnnVectorsReader knnVectorsReader = mergeState.knnVectorsReaders[i];
+        if (knnVectorsReader != null) {
+          VectorValues values = knnVectorsReader.getVectorValues(fieldInfo.name);
+          if (values != null) {
+            subs.add(new VectorValuesSub(mergeState.docMaps[i], values));
+          }
+        }
+      }
+      return new MergedVectorValues(subs, mergeState);
+    }
+
+    private MergedVectorValues(List<VectorValuesSub> subs, MergeState mergeState)
+        throws IOException {
       this.subs = subs;
       docIdMerger = DocIDMerger.of(subs, mergeState.needsIndexSort);
       int totalCost = 0, totalSize = 0;
       for (VectorValuesSub sub : subs) {
         totalCost += sub.values.cost();
         totalSize += sub.values.size();
       }
-      /* This size includes deleted docs, but when we iterate over docs here (nextDoc())
-       * we skip deleted docs. So we sneakily update this size once we observe that iteration is complete.
-       * That way by the time we are asked to do random access for graph building, we have a correct size.
-       */
       cost = totalCost;
       size = totalSize;
-      ordMap = new int[size];
-      ordBase = new int[subs.size()];
-      int lastBase = 0;
-      for (int k = 0; k < subs.size(); k++) {
-        int size = subs.get(k).values.size();
-        ordBase[k] = lastBase;
-        lastBase += size;
-      }
       docId = -1;
     }
 
@@ -227,12 +168,8 @@ public int nextDoc() throws IOException {
       current = docIdMerger.next();
       if (current == null) {
         docId = NO_MORE_DOCS;
-        /* update the size to reflect the number of *non-deleted* documents seen so we can support
-         * random access. */
-        size = ord;
       } else {
         docId = current.mappedDocID;
-        ordMap[ord++] = ordBase[current.segmentIndex] + current.count - 1;
       }
       return docId;
     }
@@ -247,11 +184,6 @@ public BytesRef binaryValue() throws IOException {
       return current.values.binaryValue();
     }
 
-    @Override
-    public RandomAccessVectorValues randomAccess() {
-      return new MergerRandomAccess();
-    }
-
     @Override
     public int advance(int target) {
       throw new UnsupportedOperationException();
@@ -271,52 +203,5 @@ public long cost() {
     public int dimension() {
       return subs.get(0).values.dimension();
     }
-
-    class MergerRandomAccess implements RandomAccessVectorValues {
-
-      private final List<RandomAccessVectorValues> raSubs;
-
-      MergerRandomAccess() {
-        raSubs = new ArrayList<>(subs.size());
-        for (VectorValuesSub sub : subs) {
-          if (sub.values instanceof RandomAccessVectorValuesProducer) {
-            raSubs.add(((RandomAccessVectorValuesProducer) sub.values).randomAccess());
-          } else {
-            throw new IllegalStateException(
-                "Cannot merge VectorValues without support for random access");
-          }
-        }
-      }
-
-      @Override
-      public int size() {
-        return size;
-      }
-
-      @Override
-      public int dimension() {
-        return VectorValuesMerger.this.dimension();
-      }
-
-      @Override
-      public float[] vectorValue(int target) throws IOException {
-        int unmappedOrd = ordMap[target];
-        int segmentOrd = Arrays.binarySearch(ordBase, unmappedOrd);
-        if (segmentOrd < 0) {
-          // get the index of the greatest lower bound
-          segmentOrd = -2 - segmentOrd;
-        }
-        while (segmentOrd < ordBase.length - 1 && ordBase[segmentOrd + 1] == ordBase[segmentOrd]) {
-          // forward over empty segments which will share the same ordBase
-          segmentOrd++;
-        }
-        return raSubs.get(segmentOrd).vectorValue(unmappedOrd - ordBase[segmentOrd]);
-      }
-
-      @Override
-      public BytesRef binaryValue(int targetOrd) throws IOException {
-        throw new UnsupportedOperationException();
-      }
-    }
   }
 }