Add match support for semantic_text fields (#117839)

* Added query name to inference field metadata

* Fix build error

* Added query builder service

* Add query builder service to query rewrite context

* Updated match query to support querying semantic text fields

* Fix build error

* Fix NPE

* Update the POC to rewrite to a bool query when combined inference and non-inference fields

* Separate clause for each inference index (to avoid inference ID clashes)

* Simplify query builder service concept to a single default inference query

* Rename QueryBuilderService, remove query name from inference metadata

* Fix too many rewrite rounds error by injecting booleans in constructors for match query builder and semantic text

* Fix test compilation errors

* Fix tests

* Add yaml test for semantic match

* Add NodeFeature

* Fix license headers

* Spotless

* Updated getClass comparison in MatchQueryBuilder

* Cleanup

* Add Mock Inference Query Builder Service

* Spotless

* Cleanup

* Update docs/changelog/117839.yaml

* Update changelog

* Replace the default inference query builder with a query rewrite interceptor

* Cleanup

* Some more cleanup/renames

* Some more cleanup/renames

* Spotless

* Checkstyle

* Convert List<QueryRewriteInterceptor> to Map keyed on query name, error on query name collisions

* PR feedback - remove check on QueryRewriteContext class only

* PR feedback

* Remove intercept flag from MatchQueryBuilder and replace with wrapper

* Move feature to test feature

* Ensure interception happens only once

* Rename InterceptedQueryBuilderWrapper to AbstractQueryBuilderWrapper

* Add lenient field to SemanticQueryBuilder

* Clean up yaml test

* Add TODO comment

* Add comment

* Spotless

* Rename AbstractQueryBuilderWrapper back to InterceptedQueryBuilderWrapper

* Spotless

* Didn't mean to commit that

* Remove static class wrapping the InterceptedQueryBuilderWrapper

* Make InterceptedQueryBuilderWrapper part of QueryRewriteInterceptor

* Refactor the interceptor to be an internal plugin that cannot be used outside inference plugin

* Fix tests

* Spotless

* Minor cleanup

* C'mon spotless

* Test spotless

* Cleanup InternalQueryRewriter

* Change if statement to assert

* Simplify template of InterceptedQueryBuilderWrapper

* Change constructor of InterceptedQueryBuilderWrapper

* Refactor InterceptedQueryBuilderWrapper to extend QueryBuilder

* Cleanup

* Add test

* Spotless

* Rename rewrite to interceptAndRewrite in QueryRewriteInterceptor

* DOESN'T WORK - for testing

* Add comment

* Getting closer - match on single typed fields works now

* Deleted line by mistake

* Checkstyle

* Fix over-aggressive IntelliJ Refactor/Rename

* And another one

* Move SemanticMatchQueryRewriteInterceptor.SEMANTIC_MATCH_QUERY_REWRITE_INTERCEPTION_SUPPORTED to Test feature

* PR feedback

* Require query name with no default

* PR feedback & update test

* Add rewrite test

* Update server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java

Co-authored-by: Mike Pellegrini <[email protected]>

---------

Co-authored-by: Mike Pellegrini <[email protected]>

Author: kderusso

docs/changelog/117839.yaml

@@ -0,0 +1,5 @@
+pr: 117839
+summary: Add match support for `semantic_text` fields
+area: "Search"
+type: enhancement
+issues: []

server/src/main/java/module-info.java

@@ -479,5 +479,5 @@
     exports org.elasticsearch.lucene.spatial;
     exports org.elasticsearch.inference.configuration;
     exports org.elasticsearch.monitor.metrics;
-
+    exports org.elasticsearch.plugins.internal.rewriter to org.elasticsearch.inference;
 }

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

@@ -137,6 +137,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion RETRIES_AND_OPERATIONS_IN_BLOBSTORE_STATS = def(8_804_00_0);
     public static final TransportVersion ADD_DATA_STREAM_OPTIONS_TO_TEMPLATES = def(8_805_00_0);
     public static final TransportVersion KNN_QUERY_RESCORE_OVERSAMPLE = def(8_806_00_0);
+    public static final TransportVersion SEMANTIC_QUERY_LENIENT = def(8_807_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -58,6 +58,7 @@
 import org.elasticsearch.indices.fielddata.cache.IndicesFieldDataCache;
 import org.elasticsearch.indices.recovery.RecoveryState;
 import org.elasticsearch.plugins.IndexStorePlugin;
+import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -478,7 +479,8 @@ public IndexService newIndexService(
         IdFieldMapper idFieldMapper,
         ValuesSourceRegistry valuesSourceRegistry,
         IndexStorePlugin.IndexFoldersDeletionListener indexFoldersDeletionListener,
-        Map<String, IndexStorePlugin.SnapshotCommitSupplier> snapshotCommitSuppliers
+        Map<String, IndexStorePlugin.SnapshotCommitSupplier> snapshotCommitSuppliers,
+        QueryRewriteInterceptor queryRewriteInterceptor
     ) throws IOException {
         final IndexEventListener eventListener = freeze();
         Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> readerWrapperFactory = indexReaderWrapper
@@ -540,7 +542,8 @@ public IndexService newIndexService(
                 indexFoldersDeletionListener,
                 snapshotCommitSupplier,
                 indexCommitListener.get(),
-                mapperMetrics
+                mapperMetrics,
+                queryRewriteInterceptor
             );
             success = true;
             return indexService;

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

@@ -85,6 +85,7 @@
 import org.elasticsearch.indices.fielddata.cache.IndicesFieldDataCache;
 import org.elasticsearch.indices.recovery.RecoveryState;
 import org.elasticsearch.plugins.IndexStorePlugin;
+import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -162,6 +163,7 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust
     private final Supplier<Sort> indexSortSupplier;
     private final ValuesSourceRegistry valuesSourceRegistry;
     private final MapperMetrics mapperMetrics;
+    private final QueryRewriteInterceptor queryRewriteInterceptor;
 
     @SuppressWarnings("this-escape")
     public IndexService(
@@ -196,7 +198,8 @@ public IndexService(
         IndexStorePlugin.IndexFoldersDeletionListener indexFoldersDeletionListener,
         IndexStorePlugin.SnapshotCommitSupplier snapshotCommitSupplier,
         Engine.IndexCommitListener indexCommitListener,
-        MapperMetrics mapperMetrics
+        MapperMetrics mapperMetrics,
+        QueryRewriteInterceptor queryRewriteInterceptor
     ) {
         super(indexSettings);
         assert indexCreationContext != IndexCreationContext.RELOAD_ANALYZERS
@@ -271,6 +274,7 @@ public IndexService(
         this.indexingOperationListeners = Collections.unmodifiableList(indexingOperationListeners);
         this.indexCommitListener = indexCommitListener;
         this.mapperMetrics = mapperMetrics;
+        this.queryRewriteInterceptor = queryRewriteInterceptor;
         try (var ignored = threadPool.getThreadContext().clearTraceContext()) {
             // kick off async ops for the first shard in this index
             this.refreshTask = new AsyncRefreshTask(this);
@@ -802,6 +806,7 @@ public QueryRewriteContext newQueryRewriteContext(
             allowExpensiveQueries,
             scriptService,
             null,
+            null,
             null
         );
     }

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

@@ -21,6 +21,7 @@
 import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.SuggestingErrorOnUnknown;
+import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
 import org.elasticsearch.xcontent.AbstractObjectParser;
 import org.elasticsearch.xcontent.FilterXContentParser;
 import org.elasticsearch.xcontent.FilterXContentParserWrapper;
@@ -278,6 +279,14 @@ protected static List<QueryBuilder> readQueries(StreamInput in) throws IOExcepti
 
     @Override
     public final QueryBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOException {
+        QueryRewriteInterceptor queryRewriteInterceptor = queryRewriteContext.getQueryRewriteInterceptor();
+        if (queryRewriteInterceptor != null) {
+            var rewritten = queryRewriteInterceptor.interceptAndRewrite(queryRewriteContext, this);
+            if (rewritten != this) {
+                return new InterceptedQueryBuilderWrapper(rewritten);
+            }
+        }
+
         QueryBuilder rewritten = doRewrite(queryRewriteContext);
         if (rewritten == this) {
             return rewritten;

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

@@ -104,6 +104,7 @@ public CoordinatorRewriteContext(
             null,
             null,
             null,
+            null,
             null
         );
         this.dateFieldRangeInfo = dateFieldRangeInfo;

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

@@ -66,6 +66,9 @@ public InnerHitBuilder innerHitBuilder() {
     public static void extractInnerHits(QueryBuilder query, Map<String, InnerHitContextBuilder> innerHitBuilders) {
         if (query instanceof AbstractQueryBuilder) {
             ((AbstractQueryBuilder<?>) query).extractInnerHitBuilders(innerHitBuilders);
+        } else if (query instanceof InterceptedQueryBuilderWrapper interceptedQuery) {
+            // Unwrap an intercepted query here
+            extractInnerHits(interceptedQuery.queryBuilder, innerHitBuilders);
         } else {
             throw new IllegalStateException(
                 "provided query builder [" + query.getClass() + "] class should inherit from AbstractQueryBuilder, but it doesn't"

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

@@ -0,0 +1,109 @@
+/*
+ * 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.index.query;
+
+import org.apache.lucene.search.Query;
+import org.elasticsearch.TransportVersion;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.Objects;
+
+/**
+ * Wrapper for instances of {@link QueryBuilder} that have been intercepted using the {@link QueryRewriteInterceptor} to
+ * break out of the rewrite phase. These instances are unwrapped on serialization.
+ */
+class InterceptedQueryBuilderWrapper implements QueryBuilder {
+
+    protected final QueryBuilder queryBuilder;
+
+    InterceptedQueryBuilderWrapper(QueryBuilder queryBuilder) {
+        super();
+        this.queryBuilder = queryBuilder;
+    }
+
+    @Override
+    public QueryBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOException {
+        QueryRewriteInterceptor queryRewriteInterceptor = queryRewriteContext.getQueryRewriteInterceptor();
+        try {
+            queryRewriteContext.setQueryRewriteInterceptor(null);
+            QueryBuilder rewritten = queryBuilder.rewrite(queryRewriteContext);
+            return rewritten != queryBuilder ? new InterceptedQueryBuilderWrapper(rewritten) : this;
+        } finally {
+            queryRewriteContext.setQueryRewriteInterceptor(queryRewriteInterceptor);
+        }
+    }
+
+    @Override
+    public String getWriteableName() {
+        return queryBuilder.getWriteableName();
+    }
+
+    @Override
+    public TransportVersion getMinimalSupportedVersion() {
+        return queryBuilder.getMinimalSupportedVersion();
+    }
+
+    @Override
+    public Query toQuery(SearchExecutionContext context) throws IOException {
+        return queryBuilder.toQuery(context);
+    }
+
+    @Override
+    public QueryBuilder queryName(String queryName) {
+        queryBuilder.queryName(queryName);
+        return this;
+    }
+
+    @Override
+    public String queryName() {
+        return queryBuilder.queryName();
+    }
+
+    @Override
+    public float boost() {
+        return queryBuilder.boost();
+    }
+
+    @Override
+    public QueryBuilder boost(float boost) {
+        queryBuilder.boost(boost);
+        return this;
+    }
+
+    @Override
+    public String getName() {
+        return queryBuilder.getName();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        queryBuilder.writeTo(out);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        return queryBuilder.toXContent(builder, params);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o instanceof InterceptedQueryBuilderWrapper == false) return false;
+        return Objects.equals(queryBuilder, ((InterceptedQueryBuilderWrapper) o).queryBuilder);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hashCode(queryBuilder);
+    }
+}

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

@@ -28,6 +28,7 @@
 import org.elasticsearch.index.mapper.MappingLookup;
 import org.elasticsearch.index.mapper.SourceFieldMapper;
 import org.elasticsearch.index.mapper.TextFieldMapper;
+import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
 import org.elasticsearch.script.ScriptCompiler;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.search.builder.PointInTimeBuilder;
@@ -70,6 +71,7 @@ public class QueryRewriteContext {
     protected Predicate<String> allowedFields;
     private final ResolvedIndices resolvedIndices;
     private final PointInTimeBuilder pit;
+    private QueryRewriteInterceptor queryRewriteInterceptor;
 
     public QueryRewriteContext(
         final XContentParserConfiguration parserConfiguration,
@@ -86,7 +88,8 @@ public QueryRewriteContext(
         final BooleanSupplier allowExpensiveQueries,
         final ScriptCompiler scriptService,
         final ResolvedIndices resolvedIndices,
-        final PointInTimeBuilder pit
+        final PointInTimeBuilder pit,
+        final QueryRewriteInterceptor queryRewriteInterceptor
     ) {
 
         this.parserConfiguration = parserConfiguration;
@@ -105,6 +108,7 @@ public QueryRewriteContext(
         this.scriptService = scriptService;
         this.resolvedIndices = resolvedIndices;
         this.pit = pit;
+        this.queryRewriteInterceptor = queryRewriteInterceptor;
     }
 
     public QueryRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis) {
@@ -123,6 +127,7 @@ public QueryRewriteContext(final XContentParserConfiguration parserConfiguration
             null,
             null,
             null,
+            null,
             null
         );
     }
@@ -132,7 +137,8 @@ public QueryRewriteContext(
         final Client client,
         final LongSupplier nowInMillis,
         final ResolvedIndices resolvedIndices,
-        final PointInTimeBuilder pit
+        final PointInTimeBuilder pit,
+        final QueryRewriteInterceptor queryRewriteInterceptor
     ) {
         this(
             parserConfiguration,
@@ -149,7 +155,8 @@ public QueryRewriteContext(
             null,
             null,
             resolvedIndices,
-            pit
+            pit,
+            queryRewriteInterceptor
         );
     }
 
@@ -428,4 +435,13 @@ public String getTierPreference() {
         // It was decided we should only test the first of these potentially multiple preferences.
         return value.split(",")[0].trim();
     }
+
+    public QueryRewriteInterceptor getQueryRewriteInterceptor() {
+        return queryRewriteInterceptor;
+    }
+
+    public void setQueryRewriteInterceptor(QueryRewriteInterceptor queryRewriteInterceptor) {
+        this.queryRewriteInterceptor = queryRewriteInterceptor;
+    }
+
 }