iter

jimczi · Dec 1, 2024 · cdae3cf · cdae3cf
1 parent 7bcd2b1
commit cdae3cf
Show file tree

Hide file tree

Showing 28 changed files with 4,889 additions and 5,768 deletions.
diff --git a/server/src/main/java/org/elasticsearch/action/update/TransportUpdateAction.java b/server/src/main/java/org/elasticsearch/action/update/TransportUpdateAction.java
@@ -42,6 +42,7 @@
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.IndexNotFoundException;
 import org.elasticsearch.index.IndexService;
+import org.elasticsearch.index.IndexVersions;
 import org.elasticsearch.index.engine.VersionConflictEngineException;
 import org.elasticsearch.index.mapper.InferenceFieldMapper;
 import org.elasticsearch.index.mapper.Mapper;
@@ -374,7 +375,8 @@ private static UpdateHelper.Result deleteInferenceResults(
         IndexMetadata indexMetadata,
         MappingLookup mappingLookup
     ) {
-        if (result.getResponseResult() != DocWriteResponse.Result.UPDATED) {
+        if (result.getResponseResult() != DocWriteResponse.Result.UPDATED
+            || indexMetadata.getCreationVersion().onOrAfter(IndexVersions.INFERENCE_METADATA_FIELDS)) {
             return result;
         }
 
@@ -403,7 +405,7 @@ private static UpdateHelper.Result deleteInferenceResults(
             String inferenceFieldName = entry.getKey();
             Mapper mapper = mappingLookup.getMapper(inferenceFieldName);
 
-            if (mapper instanceof InferenceFieldMapper inferenceFieldMapper) {
+            if (mapper instanceof InferenceFieldMapper) {
                 String[] sourceFields = entry.getValue().getSourceFields();
                 for (String sourceField : sourceFields) {
                     if (sourceField.equals(inferenceFieldName) == false

diff --git a/server/src/main/java/org/elasticsearch/index/get/ShardGetService.java b/server/src/main/java/org/elasticsearch/index/get/ShardGetService.java
@@ -10,6 +10,7 @@
 package org.elasticsearch.index.get;
 
 import org.apache.lucene.index.SortedSetDocValues;
+import org.apache.lucene.search.IndexSearcher;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.document.DocumentField;
@@ -26,6 +27,7 @@
 import org.elasticsearch.index.fieldvisitor.LeafStoredFieldLoader;
 import org.elasticsearch.index.fieldvisitor.StoredFieldLoader;
 import org.elasticsearch.index.mapper.IgnoredFieldMapper;
+import org.elasticsearch.index.mapper.InferenceMetadataFieldsMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.Mapper;
 import org.elasticsearch.index.mapper.MapperMetrics;
@@ -39,6 +41,7 @@
 import org.elasticsearch.index.shard.MultiEngineGet;
 import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.lookup.Source;
+import org.elasticsearch.xcontent.XContentType;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -288,7 +291,6 @@ private GetResult innerGetFetch(
         boolean forceSyntheticSource
     ) throws IOException {
         assert get.exists() : "method should only be called if document could be retrieved";
-
         // check first if stored fields to be loaded don't contain an object field
         MappingLookup mappingLookup = mapperService.mappingLookup();
         if (storedFields != null) {
@@ -371,6 +373,26 @@ private GetResult innerGetFetch(
             if (fetchSourceContext.hasFilter()) {
                 source = source.filter(fetchSourceContext.filter());
             }
+
+            InferenceMetadataFieldsMapper inferenceMetadata = (InferenceMetadataFieldsMapper) mappingLookup.getMapping()
+                .getMetadataMapperByName(InferenceMetadataFieldsMapper.NAME);
+            if (inferenceMetadata != null && mapperService.mappingLookup().inferenceFields().isEmpty() == false) {
+                var inferenceLoader = inferenceMetadata.fieldType()
+                    .valueFetcher(
+                        mappingLookup,
+                        mapperService.getBitSetProducer(),
+                        new IndexSearcher(docIdAndVersion.reader),
+                        XContentType.JSON
+                    );
+                inferenceLoader.setNextReader(docIdAndVersion.reader.getContext());
+                List<Object> values = inferenceLoader.fetchValues(source, docIdAndVersion.docId, List.of());
+                if (values.size() > 0) {
+                    assert values.size() == 1;
+                    var sourceMap = source.source();
+                    sourceMap.put(InferenceMetadataFieldsMapper.NAME, values.get(0));
+                    source = Source.fromMap(sourceMap, source.sourceContentType());
+                }
+            }
             sourceBytes = source.internalSourceRef();
         }
 

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/DocumentParserContext.java b/server/src/main/java/org/elasticsearch/index/mapper/DocumentParserContext.java
@@ -42,10 +42,10 @@ public abstract class DocumentParserContext {
     /**
      * Wraps a given context while allowing to override some of its behaviour by re-implementing some of the non final methods
      */
-    static class Wrapper extends DocumentParserContext {
+    private static class Wrapper extends DocumentParserContext {
         private final DocumentParserContext in;
 
-        Wrapper(ObjectMapper parent, DocumentParserContext in) {
+        private Wrapper(ObjectMapper parent, DocumentParserContext in) {
             super(parent, parent.dynamic == null ? in.dynamic : parent.dynamic, in);
             this.in = in;
         }
@@ -60,16 +60,6 @@ public boolean isWithinCopyTo() {
             return in.isWithinCopyTo();
         }
 
-        @Override
-        public boolean isWithinInferenceMetadata() {
-            return in.isWithinInferenceMetadata();
-        }
-
-        @Override
-        public void markInferenceMetadataField() {
-            in.markInferenceMetadataField();
-        }
-
         @Override
         public ContentPath path() {
             return in.path();
@@ -155,8 +145,6 @@ private enum Scope {
     // Indicates if the source for this context has been marked to be recorded. Applies to synthetic source only.
     private boolean recordedSource;
 
-    private boolean hasInferenceMetadata;
-
     private DocumentParserContext(
         MappingLookup mappingLookup,
         MappingParserContext mappingParserContext,
@@ -660,25 +648,6 @@ public boolean isWithinCopyTo() {
         return false;
     }
 
-    public boolean isWithinInferenceMetadata() {
-        return false;
-    }
-
-    /**
-     * Called by {@link InferenceMetadataFieldsMapper} to indicate whether the metadata field is present
-     * in _source.
-     */
-    public void markInferenceMetadataField() {
-        this.hasInferenceMetadata = true;
-    }
-
-    /**
-     * Returns whether the _source contains an inference metadata field.
-     */
-    public final boolean hasInferenceMetadataField() {
-        return hasInferenceMetadata;
-    }
-
     boolean inArrayScope() {
         return currentScope == Scope.ARRAY;
     }

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/InferenceMetadataFieldsMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/InferenceMetadataFieldsMapper.java
@@ -9,52 +9,20 @@
 
 package org.elasticsearch.index.mapper;
 
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
-import org.elasticsearch.common.xcontent.XContentParserUtils;
-import org.elasticsearch.index.query.QueryShardException;
-import org.elasticsearch.index.query.SearchExecutionContext;
-import org.elasticsearch.xcontent.XContentParser;
+import org.apache.lucene.search.join.BitSetProducer;
+import org.elasticsearch.xcontent.XContentType;
 
-import java.io.IOException;
 import java.util.Map;
+import java.util.function.Function;
 
-public class InferenceMetadataFieldsMapper extends MetadataFieldMapper {
+public abstract class InferenceMetadataFieldsMapper extends MetadataFieldMapper {
     public static final String NAME = "_inference_fields";
     public static final String CONTENT_TYPE = "_inference_fields";
 
-    private static final InferenceMetadataFieldsMapper INSTANCE = new InferenceMetadataFieldsMapper();
-
-    public static final TypeParser PARSER = new FixedTypeParser(c -> INSTANCE);
-
-    public static final class InferenceFieldType extends MappedFieldType {
-        private static InferenceFieldType INSTANCE = new InferenceFieldType();
-
-        public InferenceFieldType() {
-            super(NAME, false, false, false, TextSearchInfo.NONE, Map.of());
-        }
-
-        @Override
-        public ValueFetcher valueFetcher(SearchExecutionContext context, String format) {
-            // TODO: return the map from the individual semantic text fields?
-            return null;
-        }
-
-        @Override
-        public String typeName() {
-            return CONTENT_TYPE;
-        }
-
-        @Override
-        public Query termQuery(Object value, SearchExecutionContext context) {
-            throw new QueryShardException(
-                context,
-                "[" + name() + "] field which is of type [" + typeName() + "], does not support term queries"
-            );
-        }
-    }
-
-    private InferenceMetadataFieldsMapper() {
-        super(InferenceFieldType.INSTANCE);
+    protected InferenceMetadataFieldsMapper(MappedFieldType inferenceFieldType) {
+        super(inferenceFieldType);
     }
 
     @Override
@@ -63,34 +31,20 @@ protected String contentType() {
     }
 
     @Override
-    protected boolean supportsParsingObject() {
-        return true;
+    public InferenceMetadataFieldType fieldType() {
+        return (InferenceMetadataFieldType) super.fieldType();
     }
 
-    @Override
-    protected void parseCreateField(DocumentParserContext context) throws IOException {
-        XContentParser parser = context.parser();
-        XContentParserUtils.ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
-        context.markInferenceMetadataField();
-        while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
-            XContentParserUtils.ensureExpectedToken(XContentParser.Token.FIELD_NAME, parser.currentToken(), parser);
-            String fieldName = parser.currentName();
-            var parent = context.parent().findParentMapper(fieldName);
-            if (parent == null) {
-                throw new IllegalArgumentException("Illegal inference field [" + fieldName + "] found.");
-            }
-            String suffix = context.parent() != parent ? fieldName.substring(parent.fullPath().length() + 1) : fieldName;
-            var mapper = parent.getMapper(suffix);
-            if (mapper != null && mapper instanceof InferenceFieldMapper && mapper instanceof FieldMapper fieldMapper) {
-                fieldMapper.parseCreateField(new DocumentParserContext.Wrapper(context.parent(), context) {
-                    @Override
-                    public boolean isWithinInferenceMetadata() {
-                        return true;
-                    }
-                });
-            } else {
-                throw new IllegalArgumentException("Illegal inference field [" + fieldName + "] found.");
-            }
+    public abstract static class InferenceMetadataFieldType extends MappedFieldType {
+        public InferenceMetadataFieldType() {
+            super(NAME, false, false, false, TextSearchInfo.NONE, Map.of());
         }
+
+        public abstract ValueFetcher valueFetcher(
+            MappingLookup mappingLookup,
+            Function<Query, BitSetProducer> bitSetCache,
+            IndexSearcher searcher,
+            XContentType xContentType
+        );
     }
 }
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/MapperService.java b/server/src/main/java/org/elasticsearch/index/mapper/MapperService.java
@@ -171,6 +171,7 @@ public boolean isAutoUpdate() {
     private final IndexVersion indexVersionCreated;
     private final MapperRegistry mapperRegistry;
     private final Supplier<MappingParserContext> mappingParserContextSupplier;
+    private final Function<Query, BitSetProducer> bitSetProducer;
     private final MapperMetrics mapperMetrics;
 
     private volatile DocumentMapper mapper;
@@ -245,6 +246,7 @@ public MapperService(
             this::getMetadataMappers,
             this::resolveDocumentType
         );
+        this.bitSetProducer = bitSetProducer;
         this.mapperMetrics = mapperMetrics;
     }
 
@@ -819,6 +821,10 @@ public MapperRegistry getMapperRegistry() {
         return mapperRegistry;
     }
 
+    public Function<Query, BitSetProducer> getBitSetProducer() {
+        return bitSetProducer;
+    }
+
     public MapperMetrics getMapperMetrics() {
         return mapperMetrics;
     }

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/Mapping.java b/server/src/main/java/org/elasticsearch/index/mapper/Mapping.java
@@ -103,7 +103,7 @@ public <T extends MetadataFieldMapper> T getMetadataMapperByClass(Class<T> clazz
         return (T) metadataMappersMap.get(clazz);
     }
 
-    MetadataFieldMapper getMetadataMapperByName(String mapperName) {
+    public MetadataFieldMapper getMetadataMapperByName(String mapperName) {
         return metadataMappersByName.get(mapperName);
     }
 

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/ObjectMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/ObjectMapper.java
@@ -852,7 +852,7 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
 
     }
 
-    ObjectMapper findParentMapper(String leafFieldPath) {
+    public ObjectMapper findParentMapper(String leafFieldPath) {
         var pathComponents = leafFieldPath.split("\\.");
         int startPathComponent = 0;
 

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/SourceFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/SourceFieldMapper.java
@@ -399,7 +399,7 @@ public boolean isComplete() {
     }
 
     @Override
-    public void preParse(DocumentParserContext context) throws IOException {
+    public void postParse(DocumentParserContext context) throws IOException {
         BytesReference originalSource = context.sourceToParse().source();
         XContentType contentType = context.sourceToParse().getXContentType();
         final BytesReference adaptedSource = applyFilters(context, originalSource, contentType);

diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchContext.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchContext.java
@@ -26,7 +26,6 @@
 import org.elasticsearch.search.fetch.subphase.highlight.SearchHighlightContext;
 import org.elasticsearch.search.internal.ContextIndexSearcher;
 import org.elasticsearch.search.internal.SearchContext;
-import org.elasticsearch.search.internal.ShardSearchRequest;
 import org.elasticsearch.search.lookup.Source;
 import org.elasticsearch.search.rank.RankBuilder;
 import org.elasticsearch.search.rescore.RescoreContext;

diff --git a/.../src/test/java/org/elasticsearch/index/mapper/MappingLookupInferenceFieldMapperTests.java b/.../src/test/java/org/elasticsearch/index/mapper/MappingLookupInferenceFieldMapperTests.java
@@ -105,11 +105,6 @@ public InferenceFieldMetadata getMetadata(Set<String> sourcePaths) {
             return new InferenceFieldMetadata(fullPath(), INFERENCE_ID, sourcePaths.toArray(new String[0]));
         }
 
-        @Override
-        public Object getOriginalValue(Map<String, Object> sourceAsMap) {
-            return null;
-        }
-
         @Override
         protected void parseCreateField(DocumentParserContext context) throws IOException {}
 

diff --git a/x-pack/plugin/inference/build.gradle b/x-pack/plugin/inference/build.gradle
@@ -38,6 +38,7 @@ dependencies {
   testImplementation(testArtifact(project(':server')))
   testImplementation(project(':x-pack:plugin:inference:qa:test-service-plugin'))
   testImplementation project(':modules:reindex')
+  testImplementation project(':modules:mapper-extras')
   clusterPlugins project(':x-pack:plugin:inference:qa:test-service-plugin')
 
   api "com.ibm.icu:icu4j:${versions.icu4j}"

diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java
@@ -24,7 +24,6 @@
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.features.NodeFeature;
-import org.elasticsearch.index.mapper.InferenceMetadataFieldsMapper;
 import org.elasticsearch.index.mapper.Mapper;
 import org.elasticsearch.index.mapper.MetadataFieldMapper;
 import org.elasticsearch.indices.SystemIndexDescriptor;
@@ -75,6 +74,7 @@
 import org.elasticsearch.xpack.inference.mapper.OffsetSourceFieldMapper;
 import org.elasticsearch.xpack.inference.mapper.OffsetSourceMetaFieldMapper;
 import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper;
+import org.elasticsearch.xpack.inference.mapper.XPackInferenceMetadataFieldsMapper;
 import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder;
 import org.elasticsearch.xpack.inference.rank.random.RandomRankBuilder;
 import org.elasticsearch.xpack.inference.rank.random.RandomRankRetrieverBuilder;
@@ -399,8 +399,8 @@ public void close() {
     @Override
     public Map<String, MetadataFieldMapper.TypeParser> getMetadataMappers() {
         return Map.of(
-            InferenceMetadataFieldsMapper.NAME,
-            InferenceMetadataFieldsMapper.PARSER,
+            XPackInferenceMetadataFieldsMapper.NAME,
+            XPackInferenceMetadataFieldsMapper.PARSER,
             OffsetSourceMetaFieldMapper.NAME,
             OffsetSourceMetaFieldMapper.PARSER
         );