Add a simple query rewriter to support rewriting known match queries to semantic queries

Author: kderusso

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

@@ -63,6 +63,7 @@
 import org.elasticsearch.indices.recovery.RecoveryState;
 import org.elasticsearch.plugins.IndexStorePlugin;
 import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
+import org.elasticsearch.plugins.internal.rewriter.SimpleQueryRewriter;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -496,7 +497,8 @@ public IndexService newIndexService(
         ValuesSourceRegistry valuesSourceRegistry,
         IndexStorePlugin.IndexFoldersDeletionListener indexFoldersDeletionListener,
         Map<String, IndexStorePlugin.SnapshotCommitSupplier> snapshotCommitSuppliers,
-        QueryRewriteInterceptor queryRewriteInterceptor
+        QueryRewriteInterceptor queryRewriteInterceptor,
+        SimpleQueryRewriter simpleQueryRewriter
     ) throws IOException {
         final IndexEventListener eventListener = freeze();
         Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> readerWrapperFactory = indexReaderWrapper
@@ -561,6 +563,7 @@ public IndexService newIndexService(
                 indexCommitListener.get(),
                 mapperMetrics,
                 queryRewriteInterceptor,
+                simpleQueryRewriter,
                 indexingStatsSettings,
                 searchStatsSettings,
                 mergeMetrics

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

@@ -91,6 +91,7 @@
 import org.elasticsearch.indices.recovery.RecoveryState;
 import org.elasticsearch.plugins.IndexStorePlugin;
 import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
+import org.elasticsearch.plugins.internal.rewriter.SimpleQueryRewriter;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -171,6 +172,7 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust
     private final ValuesSourceRegistry valuesSourceRegistry;
     private final MapperMetrics mapperMetrics;
     private final QueryRewriteInterceptor queryRewriteInterceptor;
+    private final SimpleQueryRewriter simpleQueryRewriter;
     private final IndexingStatsSettings indexingStatsSettings;
     private final SearchStatsSettings searchStatsSettings;
     private final MergeMetrics mergeMetrics;
@@ -211,6 +213,7 @@ public IndexService(
         Engine.IndexCommitListener indexCommitListener,
         MapperMetrics mapperMetrics,
         QueryRewriteInterceptor queryRewriteInterceptor,
+        SimpleQueryRewriter simpleQueryRewriter,
         IndexingStatsSettings indexingStatsSettings,
         SearchStatsSettings searchStatsSettings,
         MergeMetrics mergeMetrics
@@ -291,6 +294,7 @@ public IndexService(
         this.indexCommitListener = indexCommitListener;
         this.mapperMetrics = mapperMetrics;
         this.queryRewriteInterceptor = queryRewriteInterceptor;
+        this.simpleQueryRewriter = simpleQueryRewriter;
         try (var ignored = threadPool.getThreadContext().clearTraceContext()) {
             // kick off async ops for the first shard in this index
             this.refreshTask = new AsyncRefreshTask(this);

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

@@ -146,6 +146,7 @@
 import org.elasticsearch.plugins.IndexStorePlugin;
 import org.elasticsearch.plugins.PluginsService;
 import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
+import org.elasticsearch.plugins.internal.rewriter.SimpleQueryRewriter;
 import org.elasticsearch.repositories.RepositoriesService;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
@@ -281,7 +282,8 @@ public class IndicesService extends AbstractLifecycleComponent
     private final PostRecoveryMerger postRecoveryMerger;
     private final List<SearchOperationListener> searchOperationListeners;
     private final QueryRewriteInterceptor queryRewriteInterceptor;
-    final SlowLogFieldProvider slowLogFieldProvider; // pkg-private for testingå
+    private final SimpleQueryRewriter simpleQueryRewriter;
+    final SlowLogFieldProvider slowLogFieldProvider; // pkg-private for testing
     private final IndexingStatsSettings indexStatsSettings;
     private final SearchStatsSettings searchStatsSettings;
     private final MergeMetrics mergeMetrics;
@@ -359,6 +361,7 @@ public void onRemoval(ShardId shardId, String fieldName, boolean wasEvicted, lon
         this.snapshotCommitSuppliers = builder.snapshotCommitSuppliers;
         this.requestCacheKeyDifferentiator = builder.requestCacheKeyDifferentiator;
         this.queryRewriteInterceptor = builder.queryRewriteInterceptor;
+        this.simpleQueryRewriter = builder.simpleQueryRewriter;
         this.mapperMetrics = builder.mapperMetrics;
         this.mergeMetrics = builder.mergeMetrics;
         // doClose() is called when shutting down a node, yet there might still be ongoing requests
@@ -834,7 +837,8 @@ private synchronized IndexService createIndexService(
             valuesSourceRegistry,
             indexFoldersDeletionListeners,
             snapshotCommitSuppliers,
-            queryRewriteInterceptor
+            queryRewriteInterceptor,
+            simpleQueryRewriter
         );
     }
 
@@ -1865,6 +1869,10 @@ public CoordinatorRewriteContextProvider getCoordinatorRewriteContextProvider(Lo
         );
     }
 
+    public SimpleQueryRewriter getSimpleQueryRewriter() {
+        return simpleQueryRewriter;
+    }
+
     /**
      * Clears the caches for the given shard id if the shard is still allocated on this node
      */

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

@@ -37,6 +37,7 @@
 import org.elasticsearch.plugins.PluginsService;
 import org.elasticsearch.plugins.internal.InternalSearchPlugin;
 import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
+import org.elasticsearch.plugins.internal.rewriter.SimpleQueryRewriter;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.search.internal.ShardSearchRequest;
@@ -83,6 +84,7 @@ public class IndicesServiceBuilder {
     MergeMetrics mergeMetrics;
     List<SearchOperationListener> searchOperationListener = List.of();
     QueryRewriteInterceptor queryRewriteInterceptor = null;
+    SimpleQueryRewriter simpleQueryRewriter = null;
     SlowLogFieldProvider slowLogFieldProvider = new SlowLogFieldProvider() {
         @Override
         public SlowLogFields create() {
@@ -301,6 +303,27 @@ public IndicesService build() {
             }));
         queryRewriteInterceptor = QueryRewriteInterceptor.multi(queryRewriteInterceptors);
 
+        var simpleQueryRewriters = pluginsService.filterPlugins(InternalSearchPlugin.class)
+            .map(InternalSearchPlugin::getSimpleQueryRewriters)
+            .flatMap(List::stream)
+            .collect(Collectors.toMap(SimpleQueryRewriter::getName, interceptor -> {
+                if (interceptor.getName() == null) {
+                    throw new IllegalArgumentException("SimpleQueryRewriter [" + interceptor.getClass().getName() + "] requires name");
+                }
+                return interceptor;
+            }, (a, b) -> {
+                throw new IllegalStateException(
+                    "Conflicting simple rewriters ["
+                        + a.getName()
+                        + "] found in ["
+                        + a.getClass().getName()
+                        + "] and ["
+                        + b.getClass().getName()
+                        + "]"
+                );
+            }));
+        simpleQueryRewriter = SimpleQueryRewriter.multi(simpleQueryRewriters);
+
         return new IndicesService(this);
     }
 }

server/src/main/java/org/elasticsearch/plugins/internal/InternalSearchPlugin.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.plugins.internal;
 
 import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
+import org.elasticsearch.plugins.internal.rewriter.SimpleQueryRewriter;
 
 import java.util.List;
 
@@ -24,4 +25,8 @@ public interface InternalSearchPlugin {
     default List<QueryRewriteInterceptor> getQueryRewriteInterceptors() {
         return emptyList();
     }
+
+    default List<SimpleQueryRewriter> getSimpleQueryRewriters() {
+        return emptyList();
+    }
 }

server/src/main/java/org/elasticsearch/plugins/internal/rewriter/SimpleQueryRewriter.java

@@ -0,0 +1,60 @@
+/*
+ * 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.plugins.internal.rewriter;
+
+import org.elasticsearch.index.query.QueryBuilder;
+
+import java.util.Map;
+
+public interface SimpleQueryRewriter {
+
+    String getName();
+
+    QueryBuilder rewrite(QueryBuilder queryBuilder);
+
+    static SimpleQueryRewriter multi(Map<String, SimpleQueryRewriter> rewriters) {
+        return rewriters.isEmpty() ? new NoOpSimpleQueryRewriter() : new CompositeSimpleQueryRewriter(rewriters);
+    }
+
+    class CompositeSimpleQueryRewriter implements SimpleQueryRewriter {
+        final String NAME = "composite";
+        private final Map<String, SimpleQueryRewriter> simpleQueryRewriters;
+
+        private CompositeSimpleQueryRewriter(Map<String, SimpleQueryRewriter> simpleQueryRewriters) {
+            this.simpleQueryRewriters = simpleQueryRewriters;
+        }
+
+        @Override
+        public String getName() {
+            return NAME;
+        }
+
+        @Override
+        public QueryBuilder rewrite(QueryBuilder queryBuilder) {
+            SimpleQueryRewriter rewriter = simpleQueryRewriters.get(queryBuilder.getName());
+            if (rewriter != null) {
+                return rewriter.rewrite(queryBuilder);
+            }
+            return queryBuilder;
+        }
+    }
+
+    class NoOpSimpleQueryRewriter implements SimpleQueryRewriter {
+        @Override
+        public QueryBuilder rewrite(QueryBuilder queryBuilder) {
+            return queryBuilder;
+        }
+
+        @Override
+        public String getName() {
+            return null;
+        }
+    }
+}

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

@@ -51,6 +51,7 @@
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.search.NestedHelper;
 import org.elasticsearch.index.shard.IndexShard;
+import org.elasticsearch.plugins.internal.rewriter.SimpleQueryRewriter;
 import org.elasticsearch.search.aggregations.SearchContextAggregations;
 import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
@@ -161,6 +162,7 @@ final class DefaultSearchContext extends SearchContext {
     private final Map<String, SearchExtBuilder> searchExtBuilders = new HashMap<>();
     private final SearchExecutionContext searchExecutionContext;
     private final FetchPhase fetchPhase;
+    private final SimpleQueryRewriter simpleQueryRewriter;
 
     DefaultSearchContext(
         ReaderContext readerContext,
@@ -174,7 +176,8 @@ final class DefaultSearchContext extends SearchContext {
         SearchService.ResultsType resultsType,
         boolean enableQueryPhaseParallelCollection,
         int minimumDocsPerSlice,
-        long memoryAccountingBufferSize
+        long memoryAccountingBufferSize,
+        SimpleQueryRewriter simpleQueryRewriter
     ) throws IOException {
         this.readerContext = readerContext;
         this.request = request;
@@ -186,6 +189,7 @@ final class DefaultSearchContext extends SearchContext {
             this.indexService = readerContext.indexService();
             this.indexShard = readerContext.indexShard();
             this.memoryAccountingBufferSize = memoryAccountingBufferSize;
+            this.simpleQueryRewriter = simpleQueryRewriter;
 
             Engine.Searcher engineSearcher = readerContext.acquireSearcher("search");
             int maximumNumberOfSlices = determineMaximumNumberOfSlices(
@@ -976,4 +980,9 @@ public IdLoader newIdLoader() {
             return IdLoader.fromLeafStoredFieldLoader();
         }
     }
+
+    @Override
+    public SimpleQueryRewriter simpleQueryRewriter() {
+        return simpleQueryRewriter;
+    }
 }

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

@@ -1402,7 +1402,8 @@ private DefaultSearchContext createSearchContext(
                 resultsType,
                 enableQueryPhaseParallelCollection,
                 minimumDocsPerSlice,
-                memoryAccountingBufferSize
+                memoryAccountingBufferSize,
+                indicesService.getSimpleQueryRewriter()
             );
             // we clone the query shard context here just for rewriting otherwise we
             // might end up with incorrect state since we are using now() or script services

server/src/main/java/org/elasticsearch/search/internal/SearchContext.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.index.query.QueryShardException;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.shard.IndexShard;
+import org.elasticsearch.plugins.internal.rewriter.SimpleQueryRewriter;
 import org.elasticsearch.search.RescoreDocIds;
 import org.elasticsearch.search.SearchExtBuilder;
 import org.elasticsearch.search.SearchShardTarget;
@@ -449,4 +450,6 @@ public String toString() {
     public abstract SourceLoader newSourceLoader(@Nullable SourceFilter sourceFilter);
 
     public abstract IdLoader newIdLoader();
+
+    public abstract SimpleQueryRewriter simpleQueryRewriter();
 }

server/src/main/java/org/elasticsearch/search/internal/SubSearchContext.java

@@ -11,6 +11,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TotalHits;
 import org.elasticsearch.index.query.ParsedQuery;
+import org.elasticsearch.plugins.internal.rewriter.SimpleQueryRewriter;
 import org.elasticsearch.search.aggregations.SearchContextAggregations;
 import org.elasticsearch.search.collapse.CollapseContext;
 import org.elasticsearch.search.fetch.FetchSearchResult;
@@ -305,4 +306,9 @@ public TotalHits getTotalHits() {
     public float getMaxScore() {
         return querySearchResult.getMaxScore();
     }
+
+    @Override
+    public SimpleQueryRewriter simpleQueryRewriter() {
+        throw new UnsupportedOperationException("Not supported");
+    }
 }