Rebuild Inference Metadata Fields During Snapshot Recovery (#120045)

This PR introduces support for reconstructing inference metadata fields that are removed from `_source` by `SourceFieldMapper#applyFilters` during operations recovery.
The inference metadata fields are retrieved using value fetchers and are re-added to `_source` under the `_inference_fields` metadata field.

Author: jimczi

server/src/main/java/org/elasticsearch/index/engine/InternalEngine.java

@@ -3161,7 +3161,7 @@ public Translog.Snapshot newChangesSnapshot(
             final Translog.Snapshot snapshot;
             if (engineConfig.getIndexSettings().isRecoverySourceSyntheticEnabled()) {
                 snapshot = new LuceneSyntheticSourceChangesSnapshot(
-                    engineConfig.getMapperService().mappingLookup(),
+                    engineConfig.getMapperService(),
                     searcher,
                     SearchBasedChangesSnapshot.DEFAULT_BATCH_SIZE,
                     maxChunkSize,
@@ -3173,6 +3173,7 @@ public Translog.Snapshot newChangesSnapshot(
                 );
             } else {
                 snapshot = new LuceneChangesSnapshot(
+                    engineConfig.getMapperService(),
                     searcher,
                     SearchBasedChangesSnapshot.DEFAULT_BATCH_SIZE,
                     fromSeqNo,

server/src/main/java/org/elasticsearch/index/engine/LuceneChangesSnapshot.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.core.Assertions;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.fieldvisitor.FieldsVisitor;
+import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.SourceFieldMapper;
 import org.elasticsearch.index.translog.Translog;
 import org.elasticsearch.transport.Transports;
@@ -46,6 +47,7 @@ public final class LuceneChangesSnapshot extends SearchBasedChangesSnapshot {
     /**
      * Creates a new "translog" snapshot from Lucene for reading operations whose seq# in the specified range.
      *
+     * @param mapperService     the mapper service for this index
      * @param engineSearcher    the internal engine searcher which will be taken over if the snapshot is opened successfully
      * @param searchBatchSize   the number of documents should be returned by each search
      * @param fromSeqNo         the min requesting seq# - inclusive
@@ -56,6 +58,7 @@ public final class LuceneChangesSnapshot extends SearchBasedChangesSnapshot {
      * @param indexVersionCreated the version on which this index was created
      */
     public LuceneChangesSnapshot(
+        MapperService mapperService,
         Engine.Searcher engineSearcher,
         int searchBatchSize,
         long fromSeqNo,
@@ -65,7 +68,7 @@ public LuceneChangesSnapshot(
         boolean accessStats,
         IndexVersion indexVersionCreated
     ) throws IOException {
-        super(engineSearcher, searchBatchSize, fromSeqNo, toSeqNo, requiredFullRange, accessStats, indexVersionCreated);
+        super(mapperService, engineSearcher, searchBatchSize, fromSeqNo, toSeqNo, requiredFullRange, accessStats, indexVersionCreated);
         this.creationThread = Assertions.ENABLED ? Thread.currentThread() : null;
         this.singleConsumer = singleConsumer;
         this.parallelArray = new ParallelArray(this.searchBatchSize);
@@ -214,20 +217,24 @@ private Translog.Operation readDocAsOp(int docIndex) throws IOException {
                 if (leaf.reader() instanceof SequentialStoredFieldsLeafReader) {
                     storedFieldsReader = ((SequentialStoredFieldsLeafReader) leaf.reader()).getSequentialStoredFieldsReader();
                     storedFieldsReaderOrd = leaf.ord;
+                    setNextSourceMetadataReader(leaf);
                 } else {
                     storedFieldsReader = null;
                     storedFieldsReaderOrd = -1;
                 }
             }
         }
+
         if (storedFieldsReader != null) {
             assert singleConsumer : "Sequential access optimization must not be enabled for multiple consumers";
             assert parallelArray.useSequentialStoredFieldsReader;
             assert storedFieldsReaderOrd == leaf.ord : storedFieldsReaderOrd + " != " + leaf.ord;
             storedFieldsReader.document(segmentDocID, fields);
         } else {
+            setNextSourceMetadataReader(leaf);
             leaf.reader().storedFields().document(segmentDocID, fields);
         }
+        final BytesReference source = fields.source() != null ? addSourceMetadata(fields.source(), segmentDocID) : null;
 
         final Translog.Operation op;
         final boolean isTombstone = parallelArray.isTombStone[docIndex];
@@ -241,7 +248,6 @@ private Translog.Operation readDocAsOp(int docIndex) throws IOException {
                 op = new Translog.Delete(id, seqNo, primaryTerm, version);
                 assert assertDocSoftDeleted(leaf.reader(), segmentDocID) : "Delete op but soft_deletes field is not set [" + op + "]";
             } else {
-                final BytesReference source = fields.source();
                 if (source == null) {
                     // TODO: Callers should ask for the range that source should be retained. Thus we should always
                     // check for the existence source once we make peer-recovery to send ops after the local checkpoint.

server/src/main/java/org/elasticsearch/index/engine/LuceneSyntheticSourceChangesSnapshot.java

@@ -13,12 +13,11 @@
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.util.ArrayUtil;
-import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.fieldvisitor.LeafStoredFieldLoader;
 import org.elasticsearch.index.fieldvisitor.StoredFieldLoader;
-import org.elasticsearch.index.mapper.MappingLookup;
+import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.SourceFieldMetrics;
 import org.elasticsearch.index.mapper.SourceLoader;
 import org.elasticsearch.index.translog.Translog;
@@ -66,7 +65,7 @@ boolean hasRecoverySourceSize() {
     private final Deque<Translog.Operation> operationQueue = new LinkedList<>();
 
     public LuceneSyntheticSourceChangesSnapshot(
-        MappingLookup mappingLookup,
+        MapperService mapperService,
         Engine.Searcher engineSearcher,
         int searchBatchSize,
         long maxMemorySizeInBytes,
@@ -76,13 +75,13 @@ public LuceneSyntheticSourceChangesSnapshot(
         boolean accessStats,
         IndexVersion indexVersionCreated
     ) throws IOException {
-        super(engineSearcher, searchBatchSize, fromSeqNo, toSeqNo, requiredFullRange, accessStats, indexVersionCreated);
+        super(mapperService, engineSearcher, searchBatchSize, fromSeqNo, toSeqNo, requiredFullRange, accessStats, indexVersionCreated);
         // a MapperService#updateMapping(...) of empty index may not have been invoked and then mappingLookup is empty
-        assert engineSearcher.getDirectoryReader().maxDoc() == 0 || mappingLookup.isSourceSynthetic()
+        assert engineSearcher.getDirectoryReader().maxDoc() == 0 || mapperService.mappingLookup().isSourceSynthetic()
             : "either an empty index or synthetic source must be enabled for proper functionality.";
         // ensure we can buffer at least one document
         this.maxMemorySizeInBytes = maxMemorySizeInBytes > 0 ? maxMemorySizeInBytes : 1;
-        this.sourceLoader = mappingLookup.newSourceLoader(null, SourceFieldMetrics.NOOP);
+        this.sourceLoader = mapperService.mappingLookup().newSourceLoader(null, SourceFieldMetrics.NOOP);
         Set<String> storedFields = sourceLoader.requiredStoredFields();
         this.storedFieldLoader = StoredFieldLoader.create(false, storedFields);
         this.lastSeenSeqNo = fromSeqNo - 1;
@@ -194,6 +193,7 @@ private Translog.Operation[] loadDocuments(List<SearchRecord> documentRecords) t
 
                 leafFieldLoader = storedFieldLoader.getLoader(leafReaderContext, null);
                 leafSourceLoader = sourceLoader.leaf(leafReaderContext.reader(), null);
+                setNextSourceMetadataReader(leafReaderContext);
             }
             int segmentDocID = docRecord.docID() - docBase;
             leafFieldLoader.advanceTo(segmentDocID);
@@ -229,17 +229,16 @@ private Translog.Operation createOperation(
                     return null;
                 }
             }
-            BytesReference source = sourceLoader.source(fieldLoader, segmentDocID).internalSourceRef();
+            var sourceBytes = addSourceMetadata(sourceLoader.source(fieldLoader, segmentDocID).internalSourceRef(), segmentDocID);
             return new Translog.Index(
                 fieldLoader.id(),
                 docRecord.seqNo(),
                 docRecord.primaryTerm(),
                 docRecord.version(),
-                source,
+                sourceBytes,
                 fieldLoader.routing(),
                 -1 // autogenerated timestamp
             );
         }
     }
-
 }

server/src/main/java/org/elasticsearch/index/engine/SearchBasedChangesSnapshot.java

@@ -22,12 +22,17 @@
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TopFieldCollectorManager;
+import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.lucene.search.Queries;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.index.IndexVersion;
+import org.elasticsearch.index.mapper.InferenceMetadataFieldsMapper;
+import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.SeqNoFieldMapper;
+import org.elasticsearch.index.mapper.ValueFetcher;
 import org.elasticsearch.index.translog.Translog;
+import org.elasticsearch.search.lookup.Source;
 
 import java.io.Closeable;
 import java.io.IOException;
@@ -44,6 +49,7 @@ public abstract class SearchBasedChangesSnapshot implements Translog.Snapshot, C
 
     private final IndexVersion indexVersionCreated;
     private final IndexSearcher indexSearcher;
+    private final ValueFetcher sourceMetadataFetcher;
     private final Closeable onClose;
 
     protected final long fromSeqNo, toSeqNo;
@@ -67,6 +73,7 @@ public abstract class SearchBasedChangesSnapshot implements Translog.Snapshot, C
      * @param indexVersionCreated  Version of the index when it was created.
      */
     protected SearchBasedChangesSnapshot(
+        MapperService mapperService,
         Engine.Searcher engineSearcher,
         int searchBatchSize,
         long fromSeqNo,
@@ -103,6 +110,19 @@ protected SearchBasedChangesSnapshot(
 
         this.accessStats = accessStats;
         this.totalHits = accessStats ? indexSearcher.count(rangeQuery(fromSeqNo, toSeqNo, indexVersionCreated)) : -1;
+        this.sourceMetadataFetcher = createSourceMetadataValueFetcher(mapperService, indexSearcher);
+    }
+
+    private ValueFetcher createSourceMetadataValueFetcher(MapperService mapperService, IndexSearcher searcher) {
+        if (mapperService.mappingLookup().inferenceFields().isEmpty()) {
+            return null;
+        }
+        var mapper = (InferenceMetadataFieldsMapper) mapperService.mappingLookup()
+            .getMapping()
+            .getMetadataMapperByName(InferenceMetadataFieldsMapper.NAME);
+        return mapper != null
+            ? mapper.fieldType().valueFetcher(mapperService.mappingLookup(), mapperService.getBitSetProducer(), searcher)
+            : null;
     }
 
     /**
@@ -184,6 +204,45 @@ protected TopDocs nextTopDocs() throws IOException {
         return results;
     }
 
+    /**
+     * Sets the reader context to enable reading metadata that was removed from the {@code _source}.
+     * This method sets up the {@code sourceMetadataFetcher} with the provided {@link LeafReaderContext},
+     * ensuring it is ready to fetch metadata for subsequent operations.
+     *
+     * <p>Note: This method should be called before {@link #addSourceMetadata(BytesReference, int)} at the start of every leaf
+     * to ensure the metadata fetcher is properly initialized.</p>
+     */
+    protected void setNextSourceMetadataReader(LeafReaderContext context) {
+        if (sourceMetadataFetcher != null) {
+            sourceMetadataFetcher.setNextReader(context);
+        }
+    }
+
+    /**
+     * Creates a new {@link Source} object by combining the provided {@code originalSource}
+     * with additional metadata fields. If the {@code sourceMetadataFetcher} is null or no metadata
+     * fields are fetched, the original source is returned unchanged.
+     *
+     * @param originalSourceBytes the original source bytes
+     * @param segmentDocID the document ID used to fetch metadata fields
+     * @return a new {@link Source} instance containing the original data and additional metadata,
+     *         or the original source if no metadata is added
+     * @throws IOException if an error occurs while fetching metadata values
+     */
+    protected BytesReference addSourceMetadata(BytesReference originalSourceBytes, int segmentDocID) throws IOException {
+        if (sourceMetadataFetcher == null) {
+            return originalSourceBytes;
+        }
+        var originalSource = Source.fromBytes(originalSourceBytes);
+        List<Object> values = sourceMetadataFetcher.fetchValues(originalSource, segmentDocID, List.of());
+        if (values.isEmpty()) {
+            return originalSourceBytes;
+        }
+        var map = originalSource.source();
+        map.put(InferenceMetadataFieldsMapper.NAME, values.get(0));
+        return Source.fromMap(map, originalSource.sourceContentType()).internalSourceRef();
+    }
+
     static IndexSearcher newIndexSearcher(Engine.Searcher engineSearcher) throws IOException {
         return new IndexSearcher(Lucene.wrapAllDocsLive(engineSearcher.getDirectoryReader()));
     }

server/src/main/java/org/elasticsearch/index/engine/TranslogDirectoryReader.java

@@ -159,18 +159,11 @@ static DirectoryReader createInMemoryReader(
             IndexWriterConfig.OpenMode.CREATE
         ).setCodec(engineConfig.getCodec());
         try (IndexWriter writer = new IndexWriter(directory, writeConfig)) {
-            writer.addDocument(parsedDocs.rootDoc());
+            writer.addDocuments(parsedDocs.docs());
             final DirectoryReader reader = open(writer);
-            if (reader.leaves().size() != 1 || reader.leaves().get(0).reader().numDocs() != 1) {
+            if (reader.leaves().size() != 1) {
                 reader.close();
-                throw new IllegalStateException(
-                    "Expected a single document segment; "
-                        + "but ["
-                        + reader.leaves().size()
-                        + " segments with "
-                        + reader.leaves().get(0).reader().numDocs()
-                        + " documents"
-                );
+                throw new IllegalStateException("Expected a single segment; " + "but  got[" + reader.leaves().size() + "] segments");
             }
             return reader;
         } catch (IOException e) {
@@ -443,7 +436,7 @@ private void readStoredFieldsDirectly(StoredFieldVisitor visitor) throws IOExcep
                 SourceFieldMapper mapper = mappingLookup.getMapping().getMetadataMapperByClass(SourceFieldMapper.class);
                 if (mapper != null) {
                     try {
-                        sourceBytes = mapper.applyFilters(mappingLookup, sourceBytes, null);
+                        sourceBytes = mapper.applyFilters(mappingLookup, sourceBytes, null, true);
                     } catch (IOException e) {
                         throw new IOException("Failed to reapply filters after reading from translog", e);
                     }

server/src/main/java/org/elasticsearch/index/engine/TranslogOperationAsserter.java

@@ -11,6 +11,7 @@
 
 import org.apache.lucene.search.similarities.BM25Similarity;
 import org.apache.lucene.store.ByteBuffersDirectory;
+import org.elasticsearch.common.lucene.index.ElasticsearchDirectoryReader;
 import org.elasticsearch.index.cache.query.TrivialQueryCachingPolicy;
 import org.elasticsearch.index.mapper.DocumentParser;
 import org.elasticsearch.index.mapper.MappingLookup;
@@ -54,7 +55,10 @@ static Translog.Index synthesizeSource(EngineConfig engineConfig, Translog.Index
         final DocumentParser documentParser = engineConfig.getMapperService().documentParser();
         try (
             var directory = new ByteBuffersDirectory();
-            var reader = TranslogDirectoryReader.createInMemoryReader(shardId, engineConfig, directory, documentParser, mappingLookup, op)
+            var reader = ElasticsearchDirectoryReader.wrap(
+                TranslogDirectoryReader.createInMemoryReader(shardId, engineConfig, directory, documentParser, mappingLookup, op),
+                new ShardId("index", "_na_", 0)
+            )
         ) {
             final Engine.Searcher searcher = new Engine.Searcher(
                 "assert_translog",
@@ -66,7 +70,7 @@ static Translog.Index synthesizeSource(EngineConfig engineConfig, Translog.Index
             );
             try (
                 LuceneSyntheticSourceChangesSnapshot snapshot = new LuceneSyntheticSourceChangesSnapshot(
-                    mappingLookup,
+                    engineConfig.getMapperService(),
                     searcher,
                     LuceneSyntheticSourceChangesSnapshot.DEFAULT_BATCH_SIZE,
                     Integer.MAX_VALUE,