Fix vector rehydration from _source during concurrent query phase (elastic#134747)

This change ensures that a separate patch loader is created for each leaf when rehydrating vectors from _source.
This is required for concurrent searches across multiple segments, since the patch loader maintains state that must be preserved between the high-level object and the leaf object.

Note: this bug only affects serverless, so it is marked as a non-issue.

Author: jimczi

server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java

@@ -2874,8 +2874,11 @@ public String toString() {
     @Override
     public SourceLoader.SyntheticVectorsLoader syntheticVectorsLoader() {
         if (isExcludeSourceVectors) {
-            var syntheticField = new IndexedSyntheticFieldLoader(indexCreatedVersion, fieldType().similarity);
-            return new SyntheticVectorsPatchFieldLoader(syntheticField, syntheticField::copyVectorAsList);
+            return new SyntheticVectorsPatchFieldLoader<>(
+                // Recreate the object for each leaf so that different segments can be searched concurrently.
+                () -> new IndexedSyntheticFieldLoader(indexCreatedVersion, fieldType().similarity),
+                IndexedSyntheticFieldLoader::copyVectorAsList
+            );
         }
         return null;
     }
@@ -2900,10 +2903,12 @@ private class IndexedSyntheticFieldLoader extends SourceLoader.DocValuesBasedSyn
 
         private final IndexVersion indexCreatedVersion;
         private final VectorSimilarity vectorSimilarity;
+        private final Thread creationThread;
 
         private IndexedSyntheticFieldLoader(IndexVersion indexCreatedVersion, VectorSimilarity vectorSimilarity) {
             this.indexCreatedVersion = indexCreatedVersion;
             this.vectorSimilarity = vectorSimilarity;
+            this.creationThread = Thread.currentThread();
         }
 
         @Override
@@ -2930,6 +2935,8 @@ private boolean shouldNormalize() {
 
         private DocValuesLoader createLoader(KnnVectorValues.DocIndexIterator iterator, boolean checkMagnitude) {
             return docId -> {
+                assert creationThread == Thread.currentThread()
+                    : "Thread mismatch: created by [" + creationThread + "], but accessed by [" + Thread.currentThread() + "]";
                 if (iterator.docID() > docId) {
                     return hasValue = false;
                 }

server/src/main/java/org/elasticsearch/index/mapper/vectors/SparseVectorFieldMapper.java

@@ -333,8 +333,11 @@ protected SyntheticSourceSupport syntheticSourceSupport() {
     @Override
     public SourceLoader.SyntheticVectorsLoader syntheticVectorsLoader() {
         if (isExcludeSourceVectors) {
-            var syntheticField = new SparseVectorSyntheticFieldLoader(fullPath(), leafName());
-            return new SyntheticVectorsPatchFieldLoader(syntheticField, syntheticField::copyAsMap);
+            return new SyntheticVectorsPatchFieldLoader<>(
+                // Recreate the object for each leaf so that different segments can be searched concurrently.
+                () -> new SparseVectorSyntheticFieldLoader(fullPath(), leafName()),
+                SparseVectorSyntheticFieldLoader::copyAsMap
+            );
         }
         return null;
     }

server/src/main/java/org/elasticsearch/index/mapper/vectors/SyntheticVectorsPatchFieldLoader.java

@@ -10,31 +10,30 @@
 package org.elasticsearch.index.mapper.vectors;
 
 import org.apache.lucene.index.LeafReaderContext;
-import org.elasticsearch.core.CheckedSupplier;
+import org.elasticsearch.core.CheckedFunction;
 import org.elasticsearch.index.mapper.SourceLoader;
 
 import java.io.IOException;
+import java.util.function.Supplier;
 
-public class SyntheticVectorsPatchFieldLoader implements SourceLoader.SyntheticVectorsLoader {
-    private final SourceLoader.SyntheticFieldLoader syntheticLoader;
-    private final CheckedSupplier<Object, IOException> copyObject;
+public class SyntheticVectorsPatchFieldLoader<T extends SourceLoader.SyntheticFieldLoader> implements SourceLoader.SyntheticVectorsLoader {
+    private final Supplier<T> syntheticLoaderSupplier;
+    private final CheckedFunction<T, Object, IOException> copyObject;
 
-    public SyntheticVectorsPatchFieldLoader(
-        SourceLoader.SyntheticFieldLoader syntheticLoader,
-        CheckedSupplier<Object, IOException> copyObject
-    ) {
-        this.syntheticLoader = syntheticLoader;
+    public SyntheticVectorsPatchFieldLoader(Supplier<T> syntheticLoaderSupplier, CheckedFunction<T, Object, IOException> copyObject) {
+        this.syntheticLoaderSupplier = syntheticLoaderSupplier;
         this.copyObject = copyObject;
     }
 
     public SourceLoader.SyntheticVectorsLoader.Leaf leaf(LeafReaderContext context) throws IOException {
+        var syntheticLoader = syntheticLoaderSupplier.get();
         var dvLoader = syntheticLoader.docValuesLoader(context.reader(), null);
         return (doc, acc) -> {
             if (dvLoader == null) {
                 return;
             }
             if (dvLoader.advanceToDoc(doc) && syntheticLoader.hasValue()) {
-                acc.add(new SourceLoader.LeafSyntheticVectorPath(syntheticLoader.fieldName(), copyObject.get()));
+                acc.add(new SourceLoader.LeafSyntheticVectorPath(syntheticLoader.fieldName(), copyObject.apply(syntheticLoader)));
             }
         };
     }

x-pack/plugin/rank-vectors/src/main/java/org/elasticsearch/xpack/rank/vectors/mapper/RankVectorsFieldMapper.java

@@ -409,8 +409,11 @@ protected SyntheticSourceSupport syntheticSourceSupport() {
     @Override
     public SourceLoader.SyntheticVectorsLoader syntheticVectorsLoader() {
         if (isExcludeSourceVectors) {
-            var syntheticField = new DocValuesSyntheticFieldLoader();
-            return new SyntheticVectorsPatchFieldLoader(syntheticField, syntheticField::copyVectorsAsList);
+            return new SyntheticVectorsPatchFieldLoader<>(
+                // Recreate the object for each leaf so that different segments can be searched concurrently.
+                DocValuesSyntheticFieldLoader::new,
+                DocValuesSyntheticFieldLoader::copyVectorsAsList
+            );
         }
         return null;
     }