Allow including semantic field embeddings in _source (#134717)

Adds support for returning `_inference_fields` (embeddings for `semantic_text` fields)
as part of `_source` when `_source.exclude_vectors` is explicitly set to `false`.
This enables use cases like reindexing documents without recomputing embeddings.
By default, embeddings remain excluded.

Author: jimczi

docs/changelog/134717.yaml

@@ -0,0 +1,5 @@
+pr: 134717
+summary: Allow including semantic field embeddings in `_source`
+area: Vector Search
+type: enhancement
+issues: []

docs/reference/elasticsearch/mapping-reference/semantic-text.md

@@ -413,6 +413,95 @@ If you want to avoid unnecessary inference and keep existing embeddings:
     * Use **partial updates through the Bulk API**.
     * Omit any `semantic_text` fields that did not change from the `doc` object in your request.
 
+## Returning semantic field embeddings in `_source`
+
+```{applies_to}
+stack: ga 9.2
+serverless: ga
+```
+
+By default, the embeddings generated for `semantic_text` fields are stored internally and **not included in `_source`** when retrieving documents.
+
+To include the full inference fields, including their embeddings, in `_source`, set the `_source.exclude_vectors` option to `false`.
+This works with the
+[Get](https://www.elastic.co/docs/api/doc/elasticsearch/operation/operation-get),
+[Search](https://www.elastic.co/docs/api/doc/elasticsearch/operation/operation-search),
+and
+[Reindex](https://www.elastic.co/docs/api/doc/elasticsearch/operation/operation-reindex)
+APIs.
+
+```console
+POST my-index/_search
+{
+  "_source": {
+    "exclude_vectors": false
+  },
+  "query": {
+    "match_all": {}
+  }
+}
+```
+
+The embeddings will appear under `_inference_fields` in `_source`.
+
+**Use cases**
+Including embeddings in `_source` is useful when you want to:
+
+* Reindex documents into another index **with the same `inference_id`** without re-running inference.
+* Export or migrate documents while preserving their embeddings.
+* Inspect or debug the raw embeddings generated for your content.
+
+### Example: Reindex while preserving embeddings
+
+```console
+POST _reindex
+{
+  "source": {
+    "index": "my-index-src",
+    "_source": {
+      "exclude_vectors": false            <1>
+    }
+  },
+  "dest": {
+    "index": "my-index-dest"
+  }
+}
+```
+
+1. Sends the source documents with their stored embeddings to the destination index.
+
+::::{warning}
+If the target index’s `semantic_text` field does **not** use the **same `inference_id`** as the source index,
+the documents will **fail the reindex task**.
+Matching `inference_id` values are required to reuse the existing embeddings.
+::::
+
+This allows documents to be re-indexed without triggering inference again, **as long as the target `semantic_text` field uses the same `inference_id` as the source**.
+
+::::{note}
+**For versions prior to 9.2.0**
+
+Older versions do not support the `exclude_vectors` option to retrieve the embeddings of the semantic text fields.
+To return the `_inference_fields`, use the `fields` option in a search request instead:
+
+```console
+POST test-index/_search
+{
+  "query": {
+    "match": {
+      "my_semantic_field": "Which country is Paris in?"
+    }
+  },
+  "fields": [
+    "_inference_fields"
+  ]
+}
+```
+
+This returns the chunked embeddings used for semantic search under `_inference_fields` in `_source`.
+Note that the `fields` option is **not** available for the Reindex API.
+::::
+
 ## Customizing `semantic_text` indexing [custom-indexing]
 
 `semantic_text` uses defaults for indexing data based on the {{infer}} endpoint

modules/reindex/src/main/java/org/elasticsearch/reindex/AbstractAsyncBulkByScrollAction.java

@@ -214,7 +214,7 @@ static <Request extends AbstractBulkByScrollRequest<Request>> SearchRequest prep
             // always include vectors in the response unless explicitly set
             var fetchSource = sourceBuilder.fetchSource();
             if (fetchSource == null) {
-                sourceBuilder.fetchSource(FetchSourceContext.FETCH_ALL_SOURCE);
+                sourceBuilder.fetchSource(FetchSourceContext.FETCH_ALL_SOURCE_EXCLUDE_INFERENCE_FIELDS);
             } else if (fetchSource.excludeVectors() == null) {
                 sourceBuilder.excludeVectors(false);
             }

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -330,6 +330,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion NEW_SEMANTIC_QUERY_INTERCEPTORS = def(9_162_0_00);
     public static final TransportVersion ESQL_LOOKUP_JOIN_ON_EXPRESSION = def(9_163_0_00);
     public static final TransportVersion INFERENCE_REQUEST_ADAPTIVE_RATE_LIMITING_REMOVED = def(9_164_0_00);
+    public static final TransportVersion SEARCH_SOURCE_EXCLUDE_INFERENCE_FIELDS_PARAM = def(9_165_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/action/bulk/TransportShardBulkAction.java

@@ -49,11 +49,9 @@
 import org.elasticsearch.index.engine.VersionConflictEngineException;
 import org.elasticsearch.index.get.GetResult;
 import org.elasticsearch.index.mapper.DocumentMapper;
-import org.elasticsearch.index.mapper.InferenceMetadataFieldsMapper;
 import org.elasticsearch.index.mapper.MapperException;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.MappingLookup;
-import org.elasticsearch.index.mapper.RoutingFieldMapper;
 import org.elasticsearch.index.mapper.SourceToParse;
 import org.elasticsearch.index.seqno.SequenceNumbers;
 import org.elasticsearch.index.shard.IndexShard;
@@ -65,6 +63,7 @@
 import org.elasticsearch.node.NodeClosedException;
 import org.elasticsearch.plugins.internal.DocumentParsingProvider;
 import org.elasticsearch.plugins.internal.XContentMeteringParserDecorator;
+import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportRequestOptions;
 import org.elasticsearch.transport.TransportService;
@@ -366,8 +365,13 @@ static boolean executeBulkItemRequest(
         if (opType == DocWriteRequest.OpType.UPDATE) {
             final UpdateRequest updateRequest = (UpdateRequest) context.getCurrent();
             try {
-                var gFields = getStoredFieldsSpec(context.getPrimary());
-                updateResult = updateHelper.prepare(updateRequest, context.getPrimary(), nowInMillisSupplier, gFields);
+                updateResult = updateHelper.prepare(
+                    updateRequest,
+                    context.getPrimary(),
+                    nowInMillisSupplier,
+                    // Include inference fields so that partial updates can still retrieve embeddings for fields that weren't updated.
+                    FetchSourceContext.FETCH_ALL_SOURCE
+                );
             } catch (Exception failure) {
                 // we may fail translating a update to index or delete operation
                 // we use index result to communicate failure while translating update request
@@ -443,16 +447,6 @@ static boolean executeBulkItemRequest(
         return true;
     }
 
-    private static String[] getStoredFieldsSpec(IndexShard indexShard) {
-        if (InferenceMetadataFieldsMapper.isEnabled(indexShard.mapperService().mappingLookup())) {
-            if (indexShard.mapperService().mappingLookup().inferenceFields().size() > 0) {
-                // Retrieves the inference metadata field containing the inference results for all semantic fields defined in the mapping.
-                return new String[] { RoutingFieldMapper.NAME, InferenceMetadataFieldsMapper.NAME };
-            }
-        }
-        return new String[] { RoutingFieldMapper.NAME };
-    }
-
     private static boolean handleMappingUpdateRequired(
         BulkPrimaryExecutionContext context,
         MappingUpdatePerformer mappingUpdater,

server/src/main/java/org/elasticsearch/action/update/TransportUpdateAction.java

@@ -54,6 +54,7 @@
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.injection.guice.Inject;
 import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.threadpool.ThreadPool.Names;
@@ -215,7 +216,14 @@ protected void shardOperation(final UpdateRequest request, final ActionListener<
             assert ThreadPool.assertCurrentThreadPool(Names.SYSTEM_WRITE, Names.WRITE);
             return deleteInferenceResults(
                 request,
-                updateHelper.prepare(request, indexShard, threadPool::absoluteTimeInMillis), // Gets the doc using the engine
+                // Gets the doc using the engine
+                updateHelper.prepare(
+                    request,
+                    indexShard,
+                    threadPool::absoluteTimeInMillis,
+                    // Exclude inference fields to ensure embeddings are recomputed.
+                    FetchSourceContext.FETCH_ALL_SOURCE_EXCLUDE_INFERENCE_FIELDS
+                ),
                 indexService.getMetadata(),
                 mappingLookup
             );

server/src/main/java/org/elasticsearch/action/update/UpdateHelper.java

@@ -33,6 +33,7 @@
 import org.elasticsearch.script.UpdateCtxMap;
 import org.elasticsearch.script.UpdateScript;
 import org.elasticsearch.script.UpsertCtxMap;
+import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.lookup.Source;
 import org.elasticsearch.search.lookup.SourceFilter;
 import org.elasticsearch.xcontent.XContentType;
@@ -58,16 +59,10 @@ public UpdateHelper(ScriptService scriptService) {
     /**
      * Prepares an update request by converting it into an index or delete request or an update response (no action).
      */
-    public Result prepare(UpdateRequest request, IndexShard indexShard, LongSupplier nowInMillis) throws IOException {
-        // TODO: Don't hard-code gFields
-        return prepare(request, indexShard, nowInMillis, new String[] { RoutingFieldMapper.NAME });
-    }
-
-    /**
-     * Prepares an update request by converting it into an index or delete request or an update response (no action).
-     */
-    public Result prepare(UpdateRequest request, IndexShard indexShard, LongSupplier nowInMillis, String[] gFields) throws IOException {
-        final GetResult getResult = indexShard.getService().getForUpdate(request.id(), request.ifSeqNo(), request.ifPrimaryTerm(), gFields);
+    public Result prepare(UpdateRequest request, IndexShard indexShard, LongSupplier nowInMillis, FetchSourceContext fetchSourceContext)
+        throws IOException {
+        final GetResult getResult = indexShard.getService()
+            .getForUpdate(request.id(), request.ifSeqNo(), request.ifPrimaryTerm(), fetchSourceContext);
         return prepare(indexShard, request, getResult, nowInMillis);
     }
 

server/src/main/java/org/elasticsearch/index/get/ShardGetService.java

@@ -58,7 +58,6 @@
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
-import java.util.stream.Collectors;
 
 import static org.elasticsearch.index.IndexSettings.INDEX_MAPPING_EXCLUDE_SOURCE_VECTORS_SETTING;
 import static org.elasticsearch.index.seqno.SequenceNumbers.UNASSIGNED_PRIMARY_TERM;
@@ -217,16 +216,16 @@ public GetResult getFromTranslog(
         );
     }
 
-    public GetResult getForUpdate(String id, long ifSeqNo, long ifPrimaryTerm, String[] gFields) throws IOException {
+    public GetResult getForUpdate(String id, long ifSeqNo, long ifPrimaryTerm, FetchSourceContext fetchSourceContext) throws IOException {
         return doGet(
             id,
-            gFields,
+            new String[] { RoutingFieldMapper.NAME },
             true,
             Versions.MATCH_ANY,
             VersionType.INTERNAL,
             ifSeqNo,
             ifPrimaryTerm,
-            FetchSourceContext.FETCH_ALL_SOURCE,
+            fetchSourceContext,
             false,
             indexShard::get
         );
@@ -290,14 +289,8 @@ private GetResult innerGetFetch(
         // check first if stored fields to be loaded don't contain an object field
         MappingLookup mappingLookup = mapperService.mappingLookup();
         final Set<String> storedFieldSet = new HashSet<>();
-        boolean hasInferenceMetadataFields = false;
         if (storedFields != null) {
             for (String field : storedFields) {
-                if (field.equals(InferenceMetadataFieldsMapper.NAME)
-                    && InferenceMetadataFieldsMapper.isEnabled(indexShard.mapperService().mappingLookup())) {
-                    hasInferenceMetadataFields = true;
-                    continue;
-                }
                 Mapper fieldMapper = mappingLookup.getMapper(field);
                 if (fieldMapper == null) {
                     if (mappingLookup.objectMappers().get(field) != null) {
@@ -313,10 +306,16 @@ private GetResult innerGetFetch(
         Map<String, DocumentField> metadataFields = null;
         DocIdAndVersion docIdAndVersion = get.docIdAndVersion();
 
-        var res = maybeExcludeSyntheticVectorFields(mappingLookup, indexSettings, fetchSourceContext, null);
+        var res = maybeExcludeVectorFields(mappingLookup, indexSettings, fetchSourceContext, null);
         if (res.v1() != fetchSourceContext) {
             fetchSourceContext = res.v1();
         }
+
+        if (mappingLookup.inferenceFields().isEmpty() == false
+            && shouldExcludeInferenceFieldsFromSource(indexSettings, fetchSourceContext) == false) {
+            storedFieldSet.add(InferenceMetadataFieldsMapper.NAME);
+        }
+
         var sourceFilter = res.v2();
         SourceLoader loader = forceSyntheticSource
             ? new SourceLoader.Synthetic(
@@ -389,7 +388,7 @@ private GetResult innerGetFetch(
                 source = source.filter(filter);
             }
 
-            if (hasInferenceMetadataFields) {
+            if (storedFieldSet.contains(InferenceMetadataFieldsMapper.NAME)) {
                 /**
                  * Adds the {@link InferenceMetadataFieldsMapper#NAME} field from the document fields
                  * to the original _source if it has been requested.
@@ -417,18 +416,29 @@ private GetResult innerGetFetch(
      * Returns {@code true} if vector fields are explicitly marked to be excluded and {@code false} otherwise.
      */
     public static boolean shouldExcludeVectorsFromSource(IndexSettings indexSettings, FetchSourceContext fetchSourceContext) {
-        if (fetchSourceContext == null || fetchSourceContext.excludeVectors() == null) {
-            return INDEX_MAPPING_EXCLUDE_SOURCE_VECTORS_SETTING.get(indexSettings.getSettings());
-        }
-        return fetchSourceContext.excludeVectors();
+        var explicit = shouldExcludeVectorsFromSourceExplicit(fetchSourceContext);
+        return explicit != null ? explicit : INDEX_MAPPING_EXCLUDE_SOURCE_VECTORS_SETTING.get(indexSettings.getSettings());
+    }
+
+    private static Boolean shouldExcludeVectorsFromSourceExplicit(FetchSourceContext fetchSourceContext) {
+        return fetchSourceContext != null ? fetchSourceContext.excludeVectors() : null;
+    }
+
+    public static boolean shouldExcludeInferenceFieldsFromSource(IndexSettings indexSettings, FetchSourceContext fetchSourceContext) {
+        var explicit = shouldExcludeInferenceFieldsFromSourceExplicit(fetchSourceContext);
+        return explicit != null ? explicit : INDEX_MAPPING_EXCLUDE_SOURCE_VECTORS_SETTING.get(indexSettings.getSettings());
+    }
+
+    private static Boolean shouldExcludeInferenceFieldsFromSourceExplicit(FetchSourceContext fetchSourceContext) {
+        return fetchSourceContext != null ? fetchSourceContext.excludeInferenceFields() : null;
     }
 
     /**
      * Returns a {@link SourceFilter} that excludes vector fields not associated with semantic text fields,
      * unless vectors are explicitly requested to be included in the source.
      * Returns {@code null} when vectors should not be filtered out.
      */
-    public static Tuple<FetchSourceContext, SourceFilter> maybeExcludeSyntheticVectorFields(
+    public static Tuple<FetchSourceContext, SourceFilter> maybeExcludeVectorFields(
         MappingLookup mappingLookup,
         IndexSettings indexSettings,
         FetchSourceContext fetchSourceContext,
@@ -457,7 +467,7 @@ public static Tuple<FetchSourceContext, SourceFilter> maybeExcludeSyntheticVecto
             }
             // Exclude vectors from semantic text fields, as they are processed separately
             return inferenceFieldsAut == null || inferenceFieldsAut.run(f.name()) == false;
-        }).map(f -> f.name()).collect(Collectors.toList());
+        }).map(MappedFieldType::name).toList();
 
         var sourceFilter = excludes.isEmpty() ? null : new SourceFilter(new String[] {}, excludes.toArray(String[]::new));
         if (lateExcludes.size() > 0) {
@@ -466,15 +476,14 @@ public static Tuple<FetchSourceContext, SourceFilter> maybeExcludeSyntheticVecto
              * This ensures that vector fields are available to sub-fetch phases, but excluded during the {@link FetchSourcePhase}.
              */
             if (fetchSourceContext != null && fetchSourceContext.excludes() != null) {
-                for (var exclude : fetchSourceContext.excludes()) {
-                    lateExcludes.add(exclude);
-                }
+                lateExcludes.addAll(Arrays.asList(fetchSourceContext.excludes()));
             }
             var newFetchSourceContext = fetchSourceContext == null
                 ? FetchSourceContext.of(true, false, null, lateExcludes.toArray(String[]::new))
                 : FetchSourceContext.of(
                     fetchSourceContext.fetchSource(),
                     fetchSourceContext.excludeVectors(),
+                    fetchSourceContext.excludeInferenceFields(),
                     fetchSourceContext.includes(),
                     lateExcludes.toArray(String[]::new)
                 );

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -945,6 +945,7 @@ public SearchSourceBuilder excludeVectors(boolean excludeVectors) {
         this.fetchSourceContext = FetchSourceContext.of(
             fetchSourceContext.fetchSource(),
             excludeVectors,
+            fetchSourceContext.excludeInferenceFields(),
             fetchSourceContext.includes(),
             fetchSourceContext.excludes()
         );