Wire query cache into sorting nested-filter computation (elastic#42906)

Don't use Lucene's default query cache when filtering in sort.

Closes elastic#42813

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

@@ -21,6 +21,8 @@
 
 import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.IndexReaderContext;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.store.AlreadyClosedException;
 import org.apache.lucene.store.Directory;
@@ -517,6 +519,13 @@ public IndexSettings getIndexSettings() {
         return indexSettings;
     }
 
+    private IndexSearcher newCachedSearcher(int shardId, IndexReaderContext context) {
+        IndexSearcher searcher = new IndexSearcher(context);
+        searcher.setQueryCache(cache().query());
+        searcher.setQueryCachingPolicy(getShard(shardId).getQueryCachingPolicy());
+        return searcher;
+    }
+
     /**
      * Creates a new QueryShardContext.
      *
@@ -525,10 +534,9 @@ public IndexSettings getIndexSettings() {
      */
     public QueryShardContext newQueryShardContext(int shardId, IndexReader indexReader, LongSupplier nowInMillis, String clusterAlias) {
         return new QueryShardContext(
-            shardId, indexSettings, indexCache.bitsetFilterCache(), indexFieldData::getForField, mapperService(),
-                similarityService(), scriptService, xContentRegistry,
-               namedWriteableRegistry, client, indexReader,
-            nowInMillis, clusterAlias);
+            shardId, indexSettings, indexCache.bitsetFilterCache(), context -> newCachedSearcher(shardId, context),
+            indexFieldData::getForField, mapperService(), similarityService(), scriptService, xContentRegistry, namedWriteableRegistry,
+            client, indexReader, nowInMillis, clusterAlias);
     }
 
     /**

server/src/main/java/org/elasticsearch/index/fielddata/IndexFieldData.java

@@ -47,6 +47,7 @@
 import org.elasticsearch.search.sort.NestedSortBuilder;
 
 import java.io.IOException;
+import java.util.function.Function;
 
 /**
  * Thread-safe utility class that allows to get per-segment values via the
@@ -114,11 +115,14 @@ public static class Nested {
             private final BitSetProducer rootFilter;
             private final Query innerQuery;
             private final NestedSortBuilder nestedSort;
+            private final Function<IndexReaderContext, IndexSearcher> searcherFactory;
 
-            public Nested(BitSetProducer rootFilter, Query innerQuery, NestedSortBuilder nestedSort) {
+            public Nested(BitSetProducer rootFilter, Query innerQuery, NestedSortBuilder nestedSort,
+                          Function<IndexReaderContext, IndexSearcher> searcherFactory) {
                 this.rootFilter = rootFilter;
                 this.innerQuery = innerQuery;
                 this.nestedSort = nestedSort;
+                this.searcherFactory = searcherFactory;
             }
 
             public Query getInnerQuery() {
@@ -143,7 +147,7 @@ public BitSet rootDocs(LeafReaderContext ctx) throws IOException {
              */
             public DocIdSetIterator innerDocs(LeafReaderContext ctx) throws IOException {
                 final IndexReaderContext topLevelCtx = ReaderUtil.getTopLevelContext(ctx);
-                IndexSearcher indexSearcher = new IndexSearcher(topLevelCtx);
+                IndexSearcher indexSearcher = searcherFactory.apply(topLevelCtx);
                 Weight weight = indexSearcher.createWeight(indexSearcher.rewrite(innerQuery), ScoreMode.COMPLETE_NO_SCORES, 1f);
                 Scorer s = weight.scorer(ctx);
                 return s == null ? null : s.iterator();

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

@@ -22,6 +22,8 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.IndexReaderContext;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.join.BitSetProducer;
 import org.apache.lucene.search.similarities.Similarity;
@@ -62,6 +64,7 @@
 import java.util.Map;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.function.LongSupplier;
 
 /**
@@ -78,6 +81,7 @@ public class QueryShardContext extends QueryRewriteContext {
     private final MapperService mapperService;
     private final SimilarityService similarityService;
     private final BitsetFilterCache bitsetFilterCache;
+    private final Function<IndexReaderContext, IndexSearcher> searcherFactory;
     private final BiFunction<MappedFieldType, String, IndexFieldData<?>> indexFieldDataService;
     private final int shardId;
     private final IndexReader reader;
@@ -91,31 +95,35 @@ public class QueryShardContext extends QueryRewriteContext {
     private NestedScope nestedScope;
 
     public QueryShardContext(int shardId, IndexSettings indexSettings, BitsetFilterCache bitsetFilterCache,
+                             Function<IndexReaderContext, IndexSearcher> searcherFactory,
                              BiFunction<MappedFieldType, String, IndexFieldData<?>> indexFieldDataLookup, MapperService mapperService,
                              SimilarityService similarityService, ScriptService scriptService, NamedXContentRegistry xContentRegistry,
                              NamedWriteableRegistry namedWriteableRegistry, Client client, IndexReader reader, LongSupplier nowInMillis,
                              String clusterAlias) {
-        this(shardId, indexSettings, bitsetFilterCache, indexFieldDataLookup, mapperService, similarityService, scriptService,
-            xContentRegistry, namedWriteableRegistry, client, reader, nowInMillis, new Index(RemoteClusterAware.buildRemoteIndexName(
-                clusterAlias, indexSettings.getIndex().getName()), indexSettings.getIndex().getUUID()));
+        this(shardId, indexSettings, bitsetFilterCache, searcherFactory, indexFieldDataLookup, mapperService, similarityService,
+            scriptService, xContentRegistry, namedWriteableRegistry, client, reader, nowInMillis,
+            new Index(RemoteClusterAware.buildRemoteIndexName(clusterAlias, indexSettings.getIndex().getName()),
+                indexSettings.getIndex().getUUID()));
     }
 
     public QueryShardContext(QueryShardContext source) {
-        this(source.shardId, source.indexSettings, source.bitsetFilterCache, source.indexFieldDataService, source.mapperService,
-                source.similarityService, source.scriptService, source.getXContentRegistry(), source.getWriteableRegistry(),
-                source.client, source.reader, source.nowInMillis, source.fullyQualifiedIndex);
+        this(source.shardId, source.indexSettings, source.bitsetFilterCache, source.searcherFactory, source.indexFieldDataService,
+                source.mapperService, source.similarityService, source.scriptService, source.getXContentRegistry(),
+                source.getWriteableRegistry(), source.client, source.reader, source.nowInMillis, source.fullyQualifiedIndex);
     }
 
     private QueryShardContext(int shardId, IndexSettings indexSettings, BitsetFilterCache bitsetFilterCache,
+                             Function<IndexReaderContext, IndexSearcher> searcherFactory,
                              BiFunction<MappedFieldType, String, IndexFieldData<?>> indexFieldDataLookup, MapperService mapperService,
                              SimilarityService similarityService, ScriptService scriptService, NamedXContentRegistry xContentRegistry,
                              NamedWriteableRegistry namedWriteableRegistry, Client client, IndexReader reader, LongSupplier nowInMillis,
                              Index fullyQualifiedIndex) {
-        super(xContentRegistry, namedWriteableRegistry,client, nowInMillis);
+        super(xContentRegistry, namedWriteableRegistry, client, nowInMillis);
         this.shardId = shardId;
         this.similarityService = similarityService;
         this.mapperService = mapperService;
         this.bitsetFilterCache = bitsetFilterCache;
+        this.searcherFactory = searcherFactory;
         this.indexFieldDataService = indexFieldDataLookup;
         this.allowUnmappedFields = indexSettings.isDefaultAllowUnmappedFields();
         this.nestedScope = new NestedScope();
@@ -160,6 +168,10 @@ public BitSetProducer bitsetFilter(Query filter) {
         return bitsetFilterCache.getBitSetProducer(filter);
     }
 
+    public IndexSearcher newCachedSearcher(IndexReaderContext context) {
+        return searcherFactory.apply(context);
+    }
+
     public <IFD extends IndexFieldData<?>> IFD getForField(MappedFieldType fieldType) {
         return (IFD) indexFieldDataService.apply(fieldType, fullyQualifiedIndex.getName());
     }

server/src/main/java/org/elasticsearch/search/sort/SortBuilder.java

@@ -189,7 +189,7 @@ protected static Nested resolveNested(QueryShardContext context, NestedSortBuild
         } else {
             parentQuery = objectMapper.nestedTypeFilter();
         }
-        return new Nested(context.bitsetFilter(parentQuery), childQuery, nestedSort);
+        return new Nested(context.bitsetFilter(parentQuery), childQuery, nestedSort, context::newCachedSearcher);
     }
 
     private static Query resolveNestedQuery(QueryShardContext context, NestedSortBuilder nestedSort, Query parentQuery) throws IOException {

server/src/test/java/org/elasticsearch/index/fielddata/AbstractFieldDataTestCase.java

@@ -163,7 +163,7 @@ public void tearDown() throws Exception {
 
     protected Nested createNested(IndexSearcher searcher, Query parentFilter, Query childFilter) throws IOException {
         BitsetFilterCache s = indexService.cache().bitsetFilterCache();
-        return new Nested(s.getBitSetProducer(parentFilter), childFilter, null);
+        return new Nested(s.getBitSetProducer(parentFilter), childFilter, null, IndexSearcher::new);
     }
 
     public void testEmpty() throws Exception {

server/src/test/java/org/elasticsearch/index/mapper/DateFieldTypeTests.java

@@ -178,7 +178,7 @@ public void testTermQuery() {
         QueryShardContext context = new QueryShardContext(0,
                 new IndexSettings(IndexMetaData.builder("foo").settings(indexSettings).build(),
                         indexSettings),
-                null, null, null, null, null, xContentRegistry(), writableRegistry(), null, null, () -> nowInMillis, null);
+                null, null, null, null, null, null, xContentRegistry(), writableRegistry(), null, null, () -> nowInMillis, null);
         MappedFieldType ft = createDefaultFieldType();
         ft.setName("field");
         String date = "2015-10-12T14:10:55";
@@ -200,7 +200,7 @@ public void testRangeQuery() throws IOException {
                 .put(IndexMetaData.SETTING_NUMBER_OF_SHARDS, 1).put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, 1).build();
         QueryShardContext context = new QueryShardContext(0,
                 new IndexSettings(IndexMetaData.builder("foo").settings(indexSettings).build(), indexSettings),
-                null, null, null, null, null, xContentRegistry(), writableRegistry(), null, null, () -> nowInMillis, null);
+                null, null, null, null, null, null, xContentRegistry(), writableRegistry(), null, null, () -> nowInMillis, null);
         MappedFieldType ft = createDefaultFieldType();
         ft.setName("field");
         String date1 = "2015-10-12T14:10:55";

server/src/test/java/org/elasticsearch/index/mapper/FieldNamesFieldTypeTests.java

@@ -66,7 +66,7 @@ public void testTermQuery() {
         when(mapperService.simpleMatchToFullName("field_name")).thenReturn(Collections.singletonList("field_name"));
 
         QueryShardContext queryShardContext = new QueryShardContext(0,
-                indexSettings, null, null, mapperService, null, null, null, null, null, null, () -> 0L, null);
+                indexSettings, null, null, null, mapperService, null, null, null, null, null, null, () -> 0L, null);
         fieldNamesFieldType.setEnabled(true);
         Query termQuery = fieldNamesFieldType.termQuery("field_name", queryShardContext);
         assertEquals(new TermQuery(new Term(FieldNamesFieldMapper.CONTENT_TYPE, "field_name")), termQuery);

server/src/test/java/org/elasticsearch/index/mapper/RangeFieldTypeTests.java

@@ -228,7 +228,7 @@ private QueryShardContext createContext() {
         Settings indexSettings = Settings.builder()
             .put(IndexMetaData.SETTING_VERSION_CREATED, Version.CURRENT).build();
         IndexSettings idxSettings = IndexSettingsModule.newIndexSettings(randomAlphaOfLengthBetween(1, 10), indexSettings);
-        return new QueryShardContext(0, idxSettings, null, null, null, null, null, xContentRegistry(),
+        return new QueryShardContext(0, idxSettings, null, null, null, null, null, null, xContentRegistry(),
             writableRegistry(), null, null, () -> nowInMillis, null);
     }
 

server/src/test/java/org/elasticsearch/index/query/IntervalQueryBuilderTests.java

@@ -364,7 +364,7 @@ public <FactoryType> FactoryType compile(Script script, ScriptContext<FactoryTyp
 
         QueryShardContext baseContext = createShardContext();
         QueryShardContext context = new QueryShardContext(baseContext.getShardId(), baseContext.getIndexSettings(),
-            null, null, baseContext.getMapperService(), null,
+            null, null, null, baseContext.getMapperService(), null,
             scriptService,
             null, null, null, null, null, null);
 

server/src/test/java/org/elasticsearch/index/query/QueryShardContextTests.java

@@ -128,7 +128,7 @@ public static QueryShardContext createQueryShardContext(String indexUuid, String
         final long nowInMillis = randomNonNegativeLong();
 
         return new QueryShardContext(
-            0, indexSettings, null, (mappedFieldType, idxName) ->
+            0, indexSettings, null, null, (mappedFieldType, idxName) ->
             mappedFieldType.fielddataBuilder(idxName).build(indexSettings, mappedFieldType, null, null, null)
             , mapperService, null, null, NamedXContentRegistry.EMPTY, new NamedWriteableRegistry(Collections.emptyList()), null, null,
             () -> nowInMillis, clusterAlias);

server/src/test/java/org/elasticsearch/index/query/RangeQueryRewriteTests.java

@@ -37,8 +37,8 @@ public class RangeQueryRewriteTests extends ESSingleNodeTestCase {
     public void testRewriteMissingField() throws Exception {
         IndexService indexService = createIndex("test");
         IndexReader reader = new MultiReader();
-        QueryRewriteContext context = new QueryShardContext(0, indexService.getIndexSettings(), null, null, indexService.mapperService(),
-                null, null, xContentRegistry(), writableRegistry(), null, reader, null, null);
+        QueryRewriteContext context = new QueryShardContext(0, indexService.getIndexSettings(), null, null, null,
+                indexService.mapperService(), null, null, xContentRegistry(), writableRegistry(), null, reader, null, null);
         RangeQueryBuilder range = new RangeQueryBuilder("foo");
         assertEquals(Relation.DISJOINT, range.getRelation(context));
     }
@@ -54,8 +54,8 @@ public void testRewriteMissingReader() throws Exception {
             .endObject().endObject());
         indexService.mapperService().merge("type",
                 new CompressedXContent(mapping), MergeReason.MAPPING_UPDATE);
-        QueryRewriteContext context = new QueryShardContext(0, indexService.getIndexSettings(), null, null, indexService.mapperService(),
-                null, null, xContentRegistry(), writableRegistry(), null, null, null, null);
+        QueryRewriteContext context = new QueryShardContext(0, indexService.getIndexSettings(), null, null, null,
+                indexService.mapperService(), null, null, xContentRegistry(), writableRegistry(), null, null, null, null);
         RangeQueryBuilder range = new RangeQueryBuilder("foo");
         // can't make assumptions on a missing reader, so it must return INTERSECT
         assertEquals(Relation.INTERSECTS, range.getRelation(context));
@@ -73,8 +73,8 @@ public void testRewriteEmptyReader() throws Exception {
         indexService.mapperService().merge("type",
                 new CompressedXContent(mapping), MergeReason.MAPPING_UPDATE);
         IndexReader reader = new MultiReader();
-        QueryRewriteContext context = new QueryShardContext(0, indexService.getIndexSettings(), null, null, indexService.mapperService(),
-                null, null, xContentRegistry(), writableRegistry(), null, reader, null, null);
+        QueryRewriteContext context = new QueryShardContext(0, indexService.getIndexSettings(), null, null, null,
+                indexService.mapperService(), null, null, xContentRegistry(), writableRegistry(), null, reader, null, null);
         RangeQueryBuilder range = new RangeQueryBuilder("foo");
         // no values -> DISJOINT
         assertEquals(Relation.DISJOINT, range.getRelation(context));

server/src/test/java/org/elasticsearch/index/search/nested/NestedSortingTests.java

@@ -792,6 +792,8 @@ public void testMultiLevelNestedSorting() throws IOException {
             assertThat(searcher.doc(topFields.scoreDocs[0].doc).get("_id"), equalTo("4"));
             assertThat(((FieldDoc) topFields.scoreDocs[0]).fields[0], equalTo(87L));
         }
+
+        searcher.getIndexReader().close();
     }
 
     private static TopFieldDocs search(QueryBuilder queryBuilder, FieldSortBuilder sortBuilder, QueryShardContext queryShardContext,

server/src/test/java/org/elasticsearch/search/aggregations/bucket/histogram/ExtendedBoundsTests.java

@@ -100,7 +100,7 @@ public void testParseAndValidate() {
         SearchContext context = mock(SearchContext.class);
         QueryShardContext qsc = new QueryShardContext(0,
                 new IndexSettings(IndexMetaData.builder("foo").settings(indexSettings).build(), indexSettings), null, null, null, null,
-                null, xContentRegistry(), writableRegistry(), null, null, () -> now, null);
+                null, null, xContentRegistry(), writableRegistry(), null, null, () -> now, null);
         when(context.getQueryShardContext()).thenReturn(qsc);
         DateFormatter formatter = DateFormatter.forPattern("dateOptionalTime");
         DocValueFormat format = new DocValueFormat.DateTime(formatter, ZoneOffset.UTC, DateFieldMapper.Resolution.MILLISECONDS);

server/src/test/java/org/elasticsearch/search/aggregations/metrics/ScriptedMetricAggregatorTests.java

@@ -407,7 +407,7 @@ protected QueryShardContext queryShardContextMock(MapperService mapperService) {
         MockScriptEngine scriptEngine = new MockScriptEngine(MockScriptEngine.NAME, SCRIPTS, Collections.emptyMap());
         Map<String, ScriptEngine> engines = Collections.singletonMap(scriptEngine.getType(), scriptEngine);
         ScriptService scriptService =  new ScriptService(Settings.EMPTY, engines, ScriptModule.CORE_CONTEXTS);
-        return new QueryShardContext(0, mapperService.getIndexSettings(), null, null, mapperService, null, scriptService,
+        return new QueryShardContext(0, mapperService.getIndexSettings(), null, null, null, mapperService, null, scriptService,
                 xContentRegistry(), writableRegistry(), null, null, System::currentTimeMillis, null);
     }
 }

server/src/test/java/org/elasticsearch/search/fetch/subphase/highlight/HighlightBuilderTests.java

@@ -277,8 +277,8 @@ public void testBuildSearchContextHighlight() throws IOException {
         Index index = new Index(randomAlphaOfLengthBetween(1, 10), "_na_");
         IndexSettings idxSettings = IndexSettingsModule.newIndexSettings(index, indexSettings);
         // shard context will only need indicesQueriesRegistry for building Query objects nested in highlighter
-        QueryShardContext mockShardContext = new QueryShardContext(0, idxSettings, null, null, null, null, null, xContentRegistry(),
-            namedWriteableRegistry, null, null, System::currentTimeMillis, null) {
+        QueryShardContext mockShardContext = new QueryShardContext(0, idxSettings, null, null, null, null, null, null,
+            xContentRegistry(), namedWriteableRegistry, null, null, System::currentTimeMillis, null) {
             @Override
             public MappedFieldType fieldMapper(String name) {
                 TextFieldMapper.Builder builder = new TextFieldMapper.Builder(name);