Improve mutability handling of Source#source()'s map

docs/changelog/141534.yaml

@@ -0,0 +1,6 @@
+area: Search
+issues:
+ - 141075
+pr: 141534
+summary: Fixes bug that prevented retrieving _inference_fields via _source_includes
+type: bug

server/src/main/java/org/elasticsearch/index/IndexFeatures.java

@@ -11,6 +11,7 @@
 
 import org.elasticsearch.features.FeatureSpecification;
 import org.elasticsearch.features.NodeFeature;
+import org.elasticsearch.index.mapper.InferenceMetadataFieldsMapper;
 
 import java.util.Set;
 
@@ -44,7 +45,8 @@ public Set<NodeFeature> getTestFeatures() {
             SYNONYMS_SET_LENIENT_ON_NON_EXISTING,
             THROW_EXCEPTION_FOR_UNKNOWN_TOKEN_IN_REST_INDEX_PUT_ALIAS_ACTION,
             THROW_EXCEPTION_ON_INDEX_CREATION_IF_UNSUPPORTED_VALUE_TYPE_IN_ALIAS,
-            SHADOWING_DIMENSIONS_AND_METRICS_IS_VALID_IN_NON_TSDB
+            SHADOWING_DIMENSIONS_AND_METRICS_IS_VALID_IN_NON_TSDB,
+            InferenceMetadataFieldsMapper.INFERENCE_FIELDS_GET_VIA_SOURCE_INCLUDES
         );
     }
 }

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

@@ -258,8 +258,7 @@ protected Source addSyntheticFields(Source originalSource, int segmentDocID) thr
         if (values.isEmpty()) {
             return originalSource;
         }
-        originalSource.source().put(InferenceMetadataFieldsMapper.NAME, values.get(0));
-        return Source.fromMap(originalSource.source(), originalSource.sourceContentType());
+        return originalSource.withMutations(map -> map.put(InferenceMetadataFieldsMapper.NAME, values.get(0)));
     }
 
     static IndexSearcher newIndexSearcher(Engine.Searcher engineSearcher) throws IOException {

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

@@ -623,9 +623,7 @@ private static Source addInferenceMetadataFields(MapperService mapperService, Le
         inferenceLoader.setNextReader(readerContext);
         List<Object> values = inferenceLoader.fetchValues(source, docId, List.of());
         if (values.size() == 1) {
-            var newSource = source.source();
-            newSource.put(InferenceMetadataFieldsMapper.NAME, values.get(0));
-            return Source.fromMap(newSource, source.sourceContentType());
+            return source.withMutations(map -> map.put(InferenceMetadataFieldsMapper.NAME, values.get(0)));
         }
         return source;
     }

server/src/main/java/org/elasticsearch/index/mapper/InferenceMetadataFieldsMapper.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.IndexVersions;
 import org.elasticsearch.index.query.SearchExecutionContext;
@@ -49,6 +50,8 @@ public abstract class InferenceMetadataFieldsMapper extends MetadataFieldMapper
     public static final String NAME = "_inference_fields";
     public static final String CONTENT_TYPE = "_inference_fields";
 
+    public static final NodeFeature INFERENCE_FIELDS_GET_VIA_SOURCE_INCLUDES = new NodeFeature("inference_fields.get_via_source_includes");
+
     protected InferenceMetadataFieldsMapper(MappedFieldType inferenceFieldType) {
         super(inferenceFieldType);
     }

server/src/main/java/org/elasticsearch/index/mapper/SourceLoader.java

@@ -24,7 +24,6 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -467,13 +466,7 @@ public void write(LeafStoredFieldLoader storedFields, int docId, XContentBuilder
      * @return a new {@link Source} with the patches applied
      */
     static Source applySyntheticVectors(Source originalSource, List<SyntheticVectorPatch> patches) {
-        Map<String, Object> newMap = originalSource.source();
-        // Make sure we have a mutable map, empty implies `Map.of()`
-        if (newMap.isEmpty()) {
-            newMap = new LinkedHashMap<>();
-        }
-        applyPatches("", newMap, patches);
-        return Source.fromMap(newMap, originalSource.sourceContentType());
+        return originalSource.withMutations(map -> applyPatches("", map, patches));
     }
 
     /**

server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchSourcePhase.java

@@ -19,7 +19,6 @@
 import org.elasticsearch.search.lookup.Source;
 import org.elasticsearch.search.lookup.SourceFilter;
 
-import java.util.HashMap;
 import java.util.Map;
 
 public final class FetchSourcePhase implements FetchSubPhase {
@@ -99,13 +98,7 @@ private Source replaceInferenceMetadataFields(SearchHit hit, Source source) {
                 if (field == null || field.getValues().isEmpty()) {
                     return source;
                 }
-                var newSource = source.source();
-                if (newSource instanceof HashMap == false) {
-                    // the map is not mutable
-                    newSource = new HashMap<>(newSource);
-                }
-                newSource.put(InferenceMetadataFieldsMapper.NAME, field.getValues().get(0));
-                return Source.fromMap(newSource, source.sourceContentType());
+                return source.withMutations(map -> map.put(InferenceMetadataFieldsMapper.NAME, field.getValues().get(0)));
             }
 
             @Override

server/src/main/java/org/elasticsearch/search/lookup/Source.java

@@ -20,13 +20,20 @@
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.HashMap;
 import java.util.Map;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 /**
- * The source of a document.  Note that, while objects returned from {@link #source()}
- * and {@link #internalSourceRef()} are immutable, objects implementing this interface
- * may not be immutable themselves.
+ * A read-only view of a document's source.
+ *
+ * <p>This interface provides access to source content as either bytes or a parsed map.
+ * A {@code Source} instance represents a stable snapshot of a particular document's
+ * data.
+ *
+ * <p>Note: The Map returned by {@link #source()} may or may not be immutable.
+ * Use {@link #withMutations(Consumer)} to modify the Map returned by {@link #source()}.
  */
 public interface Source {
 
@@ -36,9 +43,12 @@ public interface Source {
     XContentType sourceContentType();
 
     /**
-     * A map representation of the source
-     * <p>
-     * Important: This can lose precision on numbers with a decimal point. It
+     * A map representation of the source.
+     *
+     * <p><b>IMPORTANT</b>: The returned map may be immutable. To modify the source,
+     * use {@link #withMutations(java.util.function.Consumer)} instead.
+     *
+     * <p>This can lose precision on numbers with a decimal point. It
      * converts numbers like {@code "n": 1234.567} to a {@code double} which
      * only has 52 bits of precision in the mantissa. This will come up most
      * frequently when folks write nanosecond precision dates as a decimal
@@ -56,6 +66,26 @@ public interface Source {
      */
     Source filter(SourceFilter sourceFilter);
 
+    /**
+     * Returns a new Source with mutations applied to a mutable copy of the map.
+     *
+     * <p>The mutator receives a mutable {@link java.util.LinkedHashMap} copy that can be safely modified.
+     * Use this when you need to add, remove, or update fields.
+     *
+     * <pre>
+     *   Source modified = source.withMutations(map -&gt; map.put("field", "value"));
+     * </pre>
+     *
+     * @param mutator a function that modifies the map
+     * @return a new Source with the mutations applied. There is no guarantee about the mutability of the returned value's source() map.
+     */
+    default Source withMutations(java.util.function.Consumer<Map<String, Object>> mutator) {
+        Map<String, Object> sourceMap = source();
+        Map<String, Object> mutableMap = sourceMap == null ? new HashMap<>() : new HashMap<>(sourceMap);
+        mutator.accept(mutableMap);
+        return Source.fromMap(mutableMap, sourceContentType());
+    }
+
     /**
      * For the provided path, return its value in the source.
      *

server/src/test/java/org/elasticsearch/search/lookup/SourceMutabilityTests.java

@@ -0,0 +1,213 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.search.lookup;
+
+import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xcontent.XContentType;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.function.Consumer;
+
+/**
+ * Tests {@link Source#withMutations(Consumer)}. Other source mutability tests are welcomed.
+ */
+public class SourceMutabilityTests extends ESTestCase {
+
+    public void testFromEmpty() {
+        Source s1 = Source.fromMap(null, XContentType.JSON);
+        assertTrue("Should be empty", s1.source().isEmpty());
+        assertMapIsImmutable(s1.source());
+        Source s2 = s1.withMutations(map -> map.put("key", "value"));
+        assertEquals("value", s2.source().get("key"));
+
+        Source s3 = s2.withMutations(map -> map.remove("key"));
+        assertNull(s3.source().get("key"));
+        assertTrue(s3.source().isEmpty());
+        // note that the returned map does not have to be mutable
+        assertMapIsImmutable(s3.source());
+    }
+
+    public void testWithMutationsFromImmutableMap() {
+        Source original = Source.fromMap(Map.of("field1", "value1"), XContentType.JSON);
+        assertMapIsImmutable(original.source());
+
+        Source modified = original.withMutations(map -> map.put("field2", "value2"));
+
+        assertEquals(1, original.source().size());
+        assertEquals("value1", original.source().get("field1"));
+        assertNull(original.source().get("field2"));
+
+        assertEquals(2, modified.source().size());
+        assertEquals("value1", modified.source().get("field1"));
+        assertEquals("value2", modified.source().get("field2"));
+    }
+
+    public void testWithMutationsFromMutableMap() {
+        Map<String, Object> mutableMap = new LinkedHashMap<>();
+        mutableMap.put("field1", "value1");
+        Source original = Source.fromMap(mutableMap, XContentType.JSON);
+        assertMapIsMutable("fromMap(LinkedHashMap)", original.source());
+
+        Source modified = original.withMutations(map -> map.put("field2", "value2"));
+
+        assertEquals(1, original.source().size());
+        assertEquals("value1", original.source().get("field1"));
+
+        assertEquals(2, modified.source().size());
+        assertEquals("value1", modified.source().get("field1"));
+        assertEquals("value2", modified.source().get("field2"));
+    }
+
+    public void testWithMutationsFromBytes() {
+        Source original = Source.fromBytes(new BytesArray("{\"field1\":\"value1\"}"));
+        assertMapIsMutable("fromBytes", original.source());
+
+        Source modified = original.withMutations(map -> map.put("field2", "value2"));
+
+        assertEquals(1, original.source().size());
+        assertEquals("value1", original.source().get("field1"));
+
+        assertEquals(2, modified.source().size());
+        assertEquals("value1", modified.source().get("field1"));
+        assertEquals("value2", modified.source().get("field2"));
+    }
+
+    public void testWithMutationsRemoveField() {
+        Source original = Source.fromMap(Map.of("field1", "value1", "field2", "value2"), XContentType.JSON);
+
+        Source modified = original.withMutations(map -> map.remove("field1"));
+
+        // Original should be unchanged
+        assertEquals(2, original.source().size());
+        assertEquals("value1", original.source().get("field1"));
+
+        // Modified should only have field2
+        assertEquals(1, modified.source().size());
+        assertNull(modified.source().get("field1"));
+        assertEquals("value2", modified.source().get("field2"));
+    }
+
+    public void testWithMutationsUpdateField() {
+        Source original = Source.fromMap(Map.of("field1", "value1"), XContentType.JSON);
+
+        Source modified = original.withMutations(map -> map.put("field1", "updated"));
+
+        // Original should be unchanged
+        assertEquals("value1", original.source().get("field1"));
+
+        // Modified should have updated value
+        assertEquals("updated", modified.source().get("field1"));
+    }
+
+    public void testWithMutationsChaining() {
+        Source original = Source.fromMap(Map.of("field1", "value1"), XContentType.JSON);
+
+        Source modified = original.withMutations(map -> map.put("field2", "value2"))
+            .withMutations(map -> map.put("field3", "value3"))
+            .withMutations(map -> map.remove("field1"));
+
+        // Original should be unchanged
+        assertEquals(1, original.source().size());
+        assertEquals("value1", original.source().get("field1"));
+
+        // Modified should reflect all mutations
+        assertEquals(2, modified.source().size());
+        assertNull(modified.source().get("field1"));
+        assertEquals("value2", modified.source().get("field2"));
+        assertEquals("value3", modified.source().get("field3"));
+    }
+
+    public void testWithMutationsPreservesContentType() {
+        Source jsonSource = Source.fromMap(Map.of("field", "value"), XContentType.JSON);
+        Source modifiedJson = jsonSource.withMutations(map -> map.put("new", "field"));
+        assertEquals(XContentType.JSON, modifiedJson.sourceContentType());
+
+        Source smileSource = Source.fromMap(Map.of("field", "value"), XContentType.SMILE);
+        Source modifiedSmile = smileSource.withMutations(map -> map.put("new", "field"));
+        assertEquals(XContentType.SMILE, modifiedSmile.sourceContentType());
+    }
+
+    public void testWithMutationsMultipleOperationsInSingleCall() {
+        Source original = Source.fromMap(Map.of("a", "1", "b", "2"), XContentType.JSON);
+
+        Source modified = original.withMutations(map -> {
+            map.put("c", "3");
+            map.put("d", "4");
+            map.remove("a");
+            map.put("b", "updated");
+        });
+
+        // Original unchanged
+        assertEquals(2, original.source().size());
+        assertEquals("1", original.source().get("a"));
+        assertEquals("2", original.source().get("b"));
+
+        // Modified has all changes
+        assertEquals(3, modified.source().size());
+        assertNull(modified.source().get("a"));
+        assertEquals("updated", modified.source().get("b"));
+        assertEquals("3", modified.source().get("c"));
+        assertEquals("4", modified.source().get("d"));
+    }
+
+    public void testWithMutationsHandlesNullSource() {
+        Source nullSource = new NullReturningSource();
+        assertNull(nullSource.source());
+
+        Source modified = nullSource.withMutations(map -> map.put("field", "value"));
+
+        assertEquals(1, modified.source().size());
+        assertEquals("value", modified.source().get("field"));
+    }
+
+    /**
+     * A bad Source implementation that returns null from source().
+     * This is used to test that withMutations handles null gracefully.
+     */
+    private static class NullReturningSource implements Source {
+        @Override
+        public XContentType sourceContentType() {
+            return XContentType.JSON;
+        }
+
+        @Override
+        public Map<String, Object> source() {
+            return null;
+        }
+
+        @Override
+        public BytesReference internalSourceRef() {
+            return null;
+        }
+
+        @Override
+        public Source filter(SourceFilter sourceFilter) {
+            return this;
+        }
+    }
+
+    private void assertMapIsImmutable(Map<String, Object> map) {
+        expectThrows(UnsupportedOperationException.class, () -> map.put("test_key", "test_value"));
+    }
+
+    private void assertMapIsMutable(String scenario, Map<String, Object> map) {
+        try {
+            map.put("test_key", "test_value");
+            assertEquals("test put in mutable map", "test_value", map.get("test_key"));
+            map.remove("test_key");
+            assertNull("test remove from mutable map", map.get("test_key"));
+        } catch (UnsupportedOperationException e) {
+            fail(scenario + " should return a mutable map but got: " + e.getMessage());
+        }
+    }
+}