[ML] Add new inference_rescorer via ML plugin (#96960)

This is a first step in supporting models as a rescoring mechanism for second-stage retrieval.

This commit adds:

 - a new feature flag behind which this feature will reside until its feature complete
 - The ability to use a GBDT regression model (this will ultimately be changed to a new "ltr" or "search" model type I would think) as a rescorer
 - Extract indexed document fields to pass to the rescorer (this is only the first set of feature extractions planned to add)

Author: benwtrent

server/src/main/java/org/elasticsearch/index/query/DataRewriteContext.java

@@ -0,0 +1,31 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.index.query;
+
+import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.xcontent.XContentParserConfiguration;
+
+import java.util.function.LongSupplier;
+
+/**
+ * Context object used to rewrite {@link QueryBuilder} instances into simplified version on the datanode where the request is going to be
+ * executed.
+ *
+ * Note: the way search requests are executed and rewritten currently on each node is that it is done by shard. So, `DataRewriteContext`
+ * will be used in `rewrite` per shard but before the query phase.
+ */
+public class DataRewriteContext extends QueryRewriteContext {
+    public DataRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis) {
+        super(parserConfiguration, client, nowInMillis);
+    }
+
+    @Override
+    public DataRewriteContext convertToDataRewriteContext() {
+        return this;
+    }
+}

server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java

@@ -148,6 +148,13 @@ public QueryRewriteContext convertToIndexMetadataContext() {
         return mapperService != null ? this : null;
     }
 
+    /**
+     * Returns an instance of {@link DataRewriteContext} if available or null otherwise
+     */
+    public DataRewriteContext convertToDataRewriteContext() {
+        return null;
+    }
+
     /**
      * Returns the {@link MappedFieldType} for the provided field name.
      * If the field is not mapped, the behaviour depends on the index.query.parse.allow_unmapped_fields setting, which defaults to true.

server/src/main/java/org/elasticsearch/indices/IndicesService.java

@@ -100,6 +100,7 @@
 import org.elasticsearch.index.merge.MergeStats;
 import org.elasticsearch.index.query.BoolQueryBuilder;
 import org.elasticsearch.index.query.CoordinatorRewriteContextProvider;
+import org.elasticsearch.index.query.DataRewriteContext;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.recovery.RecoveryStats;
@@ -1699,6 +1700,10 @@ public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis) {
         return new QueryRewriteContext(parserConfig, client, nowInMillis);
     }
 
+    public DataRewriteContext getDataRewriteContext(LongSupplier nowInMillis) {
+        return new DataRewriteContext(parserConfig, client, nowInMillis);
+    }
+
     public CoordinatorRewriteContextProvider getCoordinatorRewriteContextProvider(LongSupplier nowInMillis) {
         return new CoordinatorRewriteContextProvider(parserConfig, client, nowInMillis, clusterService::state, this::getTimestampFieldType);
     }

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -1662,10 +1662,10 @@ private void rewriteAndFetchShardRequest(IndexShard shard, ShardSearchRequest re
                 shard.ensureShardSearchActive(b -> l.onResponse(request));
             }
         });
-        // we also do rewrite on the coordinating node (TransportSearchService) but we also need to do it here for BWC as well as
-        // AliasFilters that might need to be rewritten. These are edge-cases but we are every efficient doing the rewrite here so it's not
-        // adding a lot of overhead
-        Rewriteable.rewriteAndFetch(request.getRewriteable(), indicesService.getRewriteContext(request::nowInMillis), actionListener);
+        // we also do rewrite on the coordinating node (TransportSearchService) but we also need to do it here.
+        // AliasFilters and other things may need to be rewritten on the data node, but not per individual shard.
+        // These are uncommon-cases but we are very efficient doing the rewrite here.
+        Rewriteable.rewriteAndFetch(request.getRewriteable(), indicesService.getDataRewriteContext(request::nowInMillis), actionListener);
     }
 
     /**

test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java

@@ -46,6 +46,7 @@
 import org.elasticsearch.index.mapper.MapperRegistry;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.query.CoordinatorRewriteContext;
+import org.elasticsearch.index.query.DataRewriteContext;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.shard.IndexLongFieldRange;
@@ -332,6 +333,10 @@ protected static CoordinatorRewriteContext createCoordinatorRewriteContext(
         return serviceHolder.createCoordinatorContext(dateFieldType, min, max);
     }
 
+    protected static DataRewriteContext dataRewriteContext() {
+        return serviceHolder.createDataContext();
+    }
+
     /**
      * @return a new {@link SearchExecutionContext} based on an index with no type registered
      */
@@ -610,6 +615,10 @@ CoordinatorRewriteContext createCoordinatorContext(DateFieldMapper.DateFieldType
             );
         }
 
+        DataRewriteContext createDataContext() {
+            return new DataRewriteContext(parserConfiguration, this.client, () -> nowInMillis);
+        }
+
         ScriptModule createScriptModule(List<ScriptPlugin> scriptPlugins) {
             if (scriptPlugins == null || scriptPlugins.isEmpty()) {
                 return new ScriptModule(Settings.EMPTY, singletonList(new ScriptPlugin() {

x-pack/plugin/build.gradle

@@ -224,6 +224,7 @@ testClusters.configureEach {
   extraConfigFile serviceTokens.name, serviceTokens
 
   requiresFeature 'es.index_mode_feature_flag_registered', Version.fromString("8.0.0")
+  requiresFeature 'es.inference_rescorer_feature_flag_enabled', Version.fromString("8.10.0")
 }
 
 tasks.register('enforceApiSpecsConvention').configure {

x-pack/plugin/ml/qa/basic-multi-node/build.gradle

@@ -1,3 +1,4 @@
+import org.elasticsearch.gradle.Version
 import org.elasticsearch.gradle.internal.info.BuildParams
 
 apply plugin: 'elasticsearch.legacy-java-rest-test'
@@ -12,6 +13,7 @@ testClusters.configureEach {
   setting 'xpack.license.self_generated.type', 'trial'
   setting 'indices.lifecycle.history_index_enabled', 'false'
   setting 'slm.history_index_enabled', 'false'
+  requiresFeature 'es.inference_rescorer_feature_flag_enabled', Version.fromString("8.10.0")
 }
 
 if (BuildParams.inFipsJvm){