WIP - Semantic text query

server/src/main/java/org/elasticsearch/action/admin/indices/validate/query/TransportValidateQueryAction.java

@@ -107,7 +107,7 @@ protected void doExecute(Task task, ValidateQueryRequest request, ActionListener
         if (request.query() == null) {
             rewriteListener.onResponse(request.query());
         } else {
-            Rewriteable.rewriteAndFetch(request.query(), searchService.getRewriteContext(timeProvider), rewriteListener);
+            Rewriteable.rewriteAndFetch(request.query(), searchService.getRewriteContext(timeProvider, request), rewriteListener);
         }
     }
 

server/src/main/java/org/elasticsearch/action/explain/TransportExplainAction.java

@@ -84,7 +84,7 @@ protected void doExecute(Task task, ExplainRequest request, ActionListener<Expla
 
         assert request.query() != null;
         LongSupplier timeProvider = () -> request.nowInMillis;
-        Rewriteable.rewriteAndFetch(request.query(), searchService.getRewriteContext(timeProvider), rewriteListener);
+        Rewriteable.rewriteAndFetch(request.query(), searchService.getRewriteContext(timeProvider, request), rewriteListener);
     }
 
     @Override

server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java

@@ -452,7 +452,11 @@ void executeRequest(
                 }
             }
         });
-        Rewriteable.rewriteAndFetch(original, searchService.getRewriteContext(timeProvider::absoluteStartMillis), rewriteListener);
+        Rewriteable.rewriteAndFetch(
+            original,
+            searchService.getRewriteContext(timeProvider::absoluteStartMillis, original),
+            rewriteListener
+        );
     }
 
     static void adjustSearchType(SearchRequest searchRequest, boolean singleShard) {

server/src/main/java/org/elasticsearch/action/search/TransportSearchShardsAction.java

@@ -104,7 +104,7 @@ protected void doExecute(Task task, SearchShardsRequest searchShardsRequest, Act
         ClusterState clusterState = clusterService.state();
         Rewriteable.rewriteAndFetch(
             original,
-            searchService.getRewriteContext(timeProvider::absoluteStartMillis),
+            searchService.getRewriteContext(timeProvider::absoluteStartMillis, original),
             listener.delegateFailureAndWrap((delegate, searchRequest) -> {
                 Map<String, OriginalIndices> groupedIndices = remoteClusterService.groupIndices(
                     searchRequest.indicesOptions(),

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

@@ -725,7 +725,8 @@ public QueryRewriteContext newQueryRewriteContext(
             namedWriteableRegistry,
             valuesSourceRegistry,
             allowExpensiveQueries,
-            scriptService
+            scriptService,
+            null
         );
     }
 

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

@@ -51,6 +51,7 @@ public CoordinatorRewriteContext(
             null,
             null,
             null,
+            null,
             null
         );
         this.indexLongFieldRange = indexLongFieldRange;

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

@@ -37,6 +37,7 @@
 import java.util.function.BooleanSupplier;
 import java.util.function.LongSupplier;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 /**
  * Context object used to rewrite {@link QueryBuilder} instances into simplified version.
@@ -59,6 +60,7 @@ public class QueryRewriteContext {
     protected boolean allowUnmappedFields;
     protected boolean mapUnmappedFieldAsString;
     protected Predicate<String> allowedFields;
+    private final Map<String, Set<String>> modelsForFields;
 
     public QueryRewriteContext(
         final XContentParserConfiguration parserConfiguration,
@@ -74,7 +76,8 @@ public QueryRewriteContext(
         final NamedWriteableRegistry namedWriteableRegistry,
         final ValuesSourceRegistry valuesSourceRegistry,
         final BooleanSupplier allowExpensiveQueries,
-        final ScriptCompiler scriptService
+        final ScriptCompiler scriptService,
+        final Map<String, Set<String>> modelsForFields
     ) {
 
         this.parserConfiguration = parserConfiguration;
@@ -92,6 +95,9 @@ public QueryRewriteContext(
         this.valuesSourceRegistry = valuesSourceRegistry;
         this.allowExpensiveQueries = allowExpensiveQueries;
         this.scriptService = scriptService;
+        this.modelsForFields = modelsForFields != null ?
+            modelsForFields.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> Set.copyOf(e.getValue()))) :
+            Collections.emptyMap();
     }
 
     public QueryRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis) {
@@ -109,10 +115,36 @@ public QueryRewriteContext(final XContentParserConfiguration parserConfiguration
             null,
             null,
             null,
+            null,
             null
         );
     }
 
+    public QueryRewriteContext(
+        final XContentParserConfiguration parserConfiguration,
+        final Client client,
+        final LongSupplier nowInMillis,
+        final Map<String, Set<String>> modelsForFields
+    ) {
+        this(
+            parserConfiguration,
+            client,
+            nowInMillis,
+            null,
+            MappingLookup.EMPTY,
+            Collections.emptyMap(),
+            null,
+            null,
+            null,
+            null,
+            null,
+            null,
+            null,
+            null,
+            modelsForFields
+        );
+    }
+
     /**
      * The registry used to build new {@link XContentParser}s. Contains registered named parsers needed to parse the query.
      *
@@ -345,4 +377,9 @@ public Iterable<String> getAllFieldNames() {
             ? allFromMapping
             : () -> Iterators.concat(allFromMapping.iterator(), runtimeMappings.keySet().iterator());
     }
+
+    public Set<String> getModelsForField(String fieldName) {
+        Set<String> models = modelsForFields.get(fieldName);
+        return models != null ? models : Collections.emptySet();
+    }
 }

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

@@ -265,7 +265,8 @@ private SearchExecutionContext(
             namedWriteableRegistry,
             valuesSourceRegistry,
             allowExpensiveQueries,
-            scriptService
+            scriptService,
+            null
         );
         this.shardId = shardId;
         this.shardRequestIndex = shardRequestIndex;

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

@@ -18,6 +18,7 @@
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ResourceAlreadyExistsException;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.IndicesRequest;
 import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
 import org.elasticsearch.action.admin.indices.mapping.put.TransportAutoPutMappingAction;
 import org.elasticsearch.action.admin.indices.mapping.put.TransportPutMappingAction;
@@ -151,6 +152,7 @@
 import java.util.Collection;
 import java.util.EnumMap;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
@@ -1695,8 +1697,20 @@ public AliasFilter buildAliasFilter(ClusterState state, String index, Set<String
     /**
      * Returns a new {@link QueryRewriteContext} with the given {@code now} provider
      */
-    public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis) {
-        return new QueryRewriteContext(parserConfig, client, nowInMillis);
+    public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis, IndicesRequest indicesRequest) {
+        Index[] indices = indexNameExpressionResolver.concreteIndices(clusterService.state(), indicesRequest);
+        Map<String, Set<String>> modelsForFields = new HashMap<>();
+        for (Index index : indices) {
+            Map<String, Set<String>> fieldsForModels = indexService(index).getMetadata().getFieldsForModels();
+            for (Map.Entry<String, Set<String>> entry : fieldsForModels.entrySet()) {
+                for (String fieldName : entry.getValue()) {
+                    Set<String> models = modelsForFields.computeIfAbsent(fieldName, v -> new HashSet<>());
+                    models.add(entry.getKey());
+                }
+            }
+        }
+
+        return new QueryRewriteContext(parserConfig, client, nowInMillis, modelsForFields);
     }
 
     public DataRewriteContext getDataRewriteContext(LongSupplier nowInMillis) {

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.ElasticsearchTimeoutException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRunnable;
+import org.elasticsearch.action.IndicesRequest;
 import org.elasticsearch.action.search.CanMatchNodeRequest;
 import org.elasticsearch.action.search.CanMatchNodeResponse;
 import org.elasticsearch.action.search.SearchRequest;
@@ -1759,8 +1760,8 @@ private void rewriteAndFetchShardRequest(IndexShard shard, ShardSearchRequest re
     /**
      * Returns a new {@link QueryRewriteContext} with the given {@code now} provider
      */
-    public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis) {
-        return indicesService.getRewriteContext(nowInMillis);
+    public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis, IndicesRequest indicesRequest) {
+        return indicesService.getRewriteContext(nowInMillis, indicesRequest);
     }
 
     public CoordinatorRewriteContextProvider getCoordinatorRewriteContextProvider(LongSupplier nowInMillis) {

server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java

@@ -124,6 +124,7 @@
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.hasSize;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -1717,7 +1718,7 @@ protected void doWriteTo(StreamOutput out) throws IOException {
             NodeClient client = new NodeClient(settings, threadPool);
 
             SearchService searchService = mock(SearchService.class);
-            when(searchService.getRewriteContext(any())).thenReturn(new QueryRewriteContext(null, null, null));
+            when(searchService.getRewriteContext(any(), eq(searchRequest))).thenReturn(new QueryRewriteContext(null, null, null));
             ClusterService clusterService = new ClusterService(
                 settings,
                 new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS),

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

@@ -601,7 +601,8 @@ QueryRewriteContext createQueryRewriteContext() {
                 namedWriteableRegistry,
                 null,
                 () -> true,
-                scriptService
+                scriptService,
+                null
             );
         }
 

x-pack/plugin/inference/src/main/java/module-info.java

@@ -18,6 +18,7 @@
     requires org.apache.httpcomponents.httpcore.nio;
     requires org.apache.lucene.core;
     requires org.elasticsearch.logging;
+    requires org.apache.lucene.join;
 
     exports org.elasticsearch.xpack.inference.action;
     exports org.elasticsearch.xpack.inference.registry;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java

@@ -23,6 +23,7 @@
 import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.index.mapper.Mapper;
 import org.elasticsearch.index.mapper.MetadataFieldMapper;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.indices.SystemIndexDescriptor;
 import org.elasticsearch.inference.InferenceServiceExtension;
 import org.elasticsearch.inference.InferenceServiceRegistry;
@@ -33,6 +34,7 @@
 import org.elasticsearch.plugins.InferenceRegistryPlugin;
 import org.elasticsearch.plugins.MapperPlugin;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.plugins.SearchPlugin;
 import org.elasticsearch.plugins.SystemIndexPlugin;
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
@@ -58,6 +60,7 @@
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
 import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper;
 import org.elasticsearch.xpack.inference.mapper.SemanticTextInferenceResultFieldMapper;
+import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder;
 import org.elasticsearch.xpack.inference.registry.ModelRegistryImpl;
 import org.elasticsearch.xpack.inference.rest.RestDeleteInferenceModelAction;
 import org.elasticsearch.xpack.inference.rest.RestGetInferenceModelAction;
@@ -86,7 +89,8 @@ public class InferencePlugin extends Plugin
         ExtensiblePlugin,
         SystemIndexPlugin,
         InferenceRegistryPlugin,
-        MapperPlugin {
+        MapperPlugin,
+        SearchPlugin {
 
     /**
      * When this setting is true the verification check that
@@ -300,4 +304,13 @@ public Map<String, Mapper.TypeParser> getMappers() {
     public Map<String, MetadataFieldMapper.TypeParser> getMetadataMappers() {
         return Map.of(SemanticTextInferenceResultFieldMapper.NAME, SemanticTextInferenceResultFieldMapper.PARSER);
     }
+
+    @Override
+    public List<QuerySpec<?>> getQueries() {
+        return List.of(new QuerySpec<QueryBuilder>(
+            SemanticQueryBuilder.NAME,
+            SemanticQueryBuilder::new,
+            SemanticQueryBuilder::fromXContent
+        ));
+    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java

@@ -7,8 +7,16 @@
 
 package org.elasticsearch.xpack.inference.mapper;
 
+import org.apache.lucene.index.Term;
+import org.apache.lucene.search.BooleanClause;
+import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.BoostQuery;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TermQuery;
+import org.apache.lucene.search.join.BitSetProducer;
+import org.apache.lucene.search.join.ScoreMode;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.lucene.search.Queries;
 import org.elasticsearch.index.fielddata.FieldDataContext;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.mapper.DocumentParserContext;
@@ -21,10 +29,15 @@
 import org.elasticsearch.index.mapper.TextSearchInfo;
 import org.elasticsearch.index.mapper.ValueFetcher;
 import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.index.search.ESToParentBlockJoinQuery;
+import org.elasticsearch.inference.InferenceResults;
+import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
 
 import java.io.IOException;
 import java.util.Map;
 
+import static org.elasticsearch.action.bulk.BulkShardRequestInferenceProvider.INFERENCE_CHUNKS_RESULTS;
+
 /**
  *  A {@link FieldMapper} for semantic text fields. These fields have a model id reference, that is used for performing inference
  * at ingestion and query time.
@@ -126,5 +139,38 @@ public ValueFetcher valueFetcher(SearchExecutionContext context, String format)
         public IndexFieldData.Builder fielddataBuilder(FieldDataContext fieldDataContext) {
             throw new IllegalArgumentException("[semantic_text] fields do not support sorting, scripting or aggregating");
         }
+
+        public Query semanticQuery(
+            InferenceResults inferenceResults,
+            SearchExecutionContext context,
+            float boost,
+            String queryName
+        ) {
+            // Cant use QueryBuilders.boolQuery() because a mapper is not registered for <field>.inference, causing
+            // TermQueryBuilder#doToQuery to fail (at TermQueryBuilder:202)
+            // TODO: Handle boost and queryName
+            String fieldName = name() + "." + INFERENCE_CHUNKS_RESULTS;
+            BooleanQuery.Builder queryBuilder = new BooleanQuery.Builder().setMinimumNumberShouldMatch(1);
+
+            // TODO: Support dense vectors
+            if (inferenceResults instanceof TextExpansionResults textExpansionResults) {
+                for (TextExpansionResults.WeightedToken weightedToken : textExpansionResults.getWeightedTokens()) {
+                    queryBuilder.add(
+                        new BoostQuery(
+                            new TermQuery(
+                                new Term(fieldName, weightedToken.token())
+                            ),
+                            weightedToken.weight()
+                        ),
+                        BooleanClause.Occur.SHOULD
+                    );
+                }
+            } else {
+                throw new IllegalArgumentException("Unsupported inference results type [" + inferenceResults.getWriteableName() + "]");
+            }
+
+            BitSetProducer parentFilter = context.bitsetFilter(Queries.newNonNestedFilter(context.indexVersionCreated()));
+            return new ESToParentBlockJoinQuery(queryBuilder.build(), parentFilter, ScoreMode.Total, name());
+        }
     }
 }