[9.2][DiskBBQ] Fix index sorting on flush (2nd attempt)

docs/changelog/145076.yaml

@@ -0,0 +1,5 @@
+area: Vector Search
+issues: []
+pr: 145076
+summary: "[DiskBBQ] Fix index sorting on flush (2nd attempt)"
+type: bug

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/IVFVectorsWriter.java

@@ -14,6 +14,7 @@
 import org.apache.lucene.codecs.KnnVectorsWriter;
 import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
+import org.apache.lucene.index.DocsWithFieldSet;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.IndexFileNames;
@@ -222,39 +223,76 @@ private static FloatVectorValues getFloatVectorValues(
         List<float[]> vectors = fieldVectorsWriter.getVectors();
         if (vectors.size() == maxDoc && sortMap == null) {
             return FloatVectorValues.fromFloats(vectors, fieldInfo.getVectorDimension());
-        }
-        final DocIdSetIterator iterator = fieldVectorsWriter.getDocsWithFieldSet().iterator();
-        final int[] docIds = new int[vectors.size()];
-        for (int i = 0; i < docIds.length; i++) {
-            docIds[i] = sortMap == null ? iterator.nextDoc() : sortMap.oldToNew(iterator.nextDoc());
-        }
-        assert iterator.nextDoc() == NO_MORE_DOCS;
-        return new FloatVectorValues() {
-            @Override
-            public float[] vectorValue(int ord) {
-                return vectors.get(ord);
+        } else if (sortMap == null) {
+            final DocIdSetIterator iterator = fieldVectorsWriter.getDocsWithFieldSet().iterator();
+            final int[] docIds = new int[vectors.size()];
+            for (int i = 0; i < docIds.length; i++) {
+                docIds[i] = iterator.nextDoc();
             }
+            assert iterator.nextDoc() == NO_MORE_DOCS;
+            return new FloatVectorValues() {
+                @Override
+                public float[] vectorValue(int ord) {
+                    return vectors.get(ord);
+                }
 
-            @Override
-            public FloatVectorValues copy() {
-                return this;
-            }
+                @Override
+                public FloatVectorValues copy() {
+                    return this;
+                }
 
-            @Override
-            public int dimension() {
-                return fieldInfo.getVectorDimension();
-            }
+                @Override
+                public int dimension() {
+                    return fieldInfo.getVectorDimension();
+                }
 
-            @Override
-            public int size() {
-                return vectors.size();
-            }
+                @Override
+                public int size() {
+                    return vectors.size();
+                }
 
-            @Override
-            public int ordToDoc(int ord) {
-                return docIds[ord];
+                @Override
+                public int ordToDoc(int ord) {
+                    return docIds[ord];
+                }
+            };
+        } else {
+            DocsWithFieldSet newDocsWithField = new DocsWithFieldSet();
+            final int[] ordMap = new int[fieldVectorsWriter.getDocsWithFieldSet().cardinality()]; // new ord to old ord
+            KnnVectorsWriter.mapOldOrdToNewOrd(fieldVectorsWriter.getDocsWithFieldSet(), sortMap, null, ordMap, newDocsWithField);
+            final DocIdSetIterator iterator = newDocsWithField.iterator();
+            final int[] docIds = new int[vectors.size()];
+            for (int i = 0; i < docIds.length; i++) {
+                docIds[i] = iterator.nextDoc();
             }
-        };
+            assert iterator.nextDoc() == NO_MORE_DOCS;
+            return new FloatVectorValues() {
+                @Override
+                public float[] vectorValue(int ord) {
+                    return vectors.get(ordMap[ord]);
+                }
+
+                @Override
+                public FloatVectorValues copy() {
+                    return this;
+                }
+
+                @Override
+                public int dimension() {
+                    return fieldInfo.getVectorDimension();
+                }
+
+                @Override
+                public int size() {
+                    return vectors.size();
+                }
+
+                @Override
+                public int ordToDoc(int ord) {
+                    return docIds[ord];
+                }
+            };
+        }
     }
 
     @Override