Minor modifications to Lucene95HnswVectorsWriter

Signed-off-by: John Mazanec <jmazane@amazon.com>
apache · Jan 30, 2023 · 4cc64ae · 4cc64ae
1 parent ca6861f
commit 4cc64ae
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene95/Lucene95HnswVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene95/Lucene95HnswVectorsWriter.java
@@ -559,8 +559,12 @@ private HnswGraph getHnswGraphFromReader(String fieldName, KnnVectorsReader knnV
       return ((Lucene95HnswVectorsReader) knnVectorsReader).getGraph(fieldName);
     }
 
+    // We should not reach here because knnVectorsReader's type is checked in
+    // selectGraphForInitialization
     throw new IllegalArgumentException(
-        "Invalid KnnVectorsReader. Must be of type PerFieldKnnVectorsFormat.FieldsReader or Lucene94HnswVectorsReader");
+        "Invalid KnnVectorsReader type for field: "
+            + fieldName
+            + ". Must be Lucene95HnswVectorsReader or newer");
   }
 
   private Map<Integer, Integer> getOldToNewOrdinalMap(
@@ -571,22 +575,26 @@ private Map<Integer, Integer> getOldToNewOrdinalMap(
 
     Map<Integer, Integer> newIdToOldOrdinal = new HashMap<>();
     int oldOrd = 0;
+    int maxNewDocID = -1;
     for (int oldId = initializerVectorValues.nextDoc();
         oldId != NO_MORE_DOCS;
         oldId = initializerVectorValues.nextDoc()) {
       if (initializerVectorValues.vectorValue() == null) {
         continue;
       }
       int newId = initializerDocMap.get(oldId);
+      maxNewDocID = Math.max(newId, maxNewDocID);
       newIdToOldOrdinal.put(newId, oldOrd);
       oldOrd++;
     }
 
+    if (maxNewDocID == -1) {
+      return Collections.emptyMap();
+    }
+
     Map<Integer, Integer> oldToNewOrdinalMap = new HashMap<>();
-    int newOrd = 0;
-    int maxNewDocID = Collections.max(newIdToOldOrdinal.keySet());
     VectorValues vectorValues = MergedVectorValues.mergeVectorValues(fieldInfo, mergeState);
-
+    int newOrd = 0;
     for (int newDocId = vectorValues.nextDoc();
         newDocId <= maxNewDocID;
         newDocId = vectorValues.nextDoc()) {