From 666676f065843e921edb0281539adf8f4899c992 Mon Sep 17 00:00:00 2001 From: Luca Cavanna Date: Wed, 19 Apr 2023 16:29:43 +0200 Subject: [PATCH 1/7] Double check if creating the collectors first is still needed --- .../search/query/QueryPhase.java | 20 +++++++------------ 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java index ffba1a0107d38..4db04bdfdd5f0 100644 --- a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java @@ -126,8 +126,12 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut } final LinkedList collectors = new LinkedList<>(); - // whether the chain contains a collector that filters documents - boolean hasFilterCollector = false; + + // create the top docs collector last when the other collectors are known + final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, + searchContext.parsedPostFilter() != null || searchContext.minimumScore() != null); + collectors.add(topDocsFactory); + if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) { // add terminate_after before the filter collectors // it will only be applied on documents accepted by these filter collectors @@ -137,8 +141,6 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut // add post filters before aggregations // it will only be applied to top hits collectors.add(createFilteredCollectorContext(searcher, searchContext.parsedPostFilter().query())); - // this collector can filter documents during the collection - hasFilterCollector = true; } if (searchContext.getAggsCollector() != null) { // plug in additional collectors, like aggregations @@ -147,8 +149,6 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut if (searchContext.minimumScore() != null) { // apply the minimum score after multi collector so we filter aggs as well collectors.add(createMinScoreCollectorContext(searchContext.minimumScore())); - // this collector can filter documents during the collection - hasFilterCollector = true; } boolean timeoutSet = scrollContext == null @@ -171,7 +171,7 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut } try { - searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, timeoutSet); + searchWithCollector(searchContext, searcher, query, collectors, timeoutSet); ExecutorService executor = searchContext.indexShard().getThreadPool().executor(ThreadPool.Names.SEARCH); assert executor instanceof EWMATrackingEsThreadPoolExecutor || (executor instanceof EsThreadPoolExecutor == false /* in case thread pool is mocked out in tests */) @@ -197,14 +197,8 @@ private static void searchWithCollector( ContextIndexSearcher searcher, Query query, LinkedList collectors, - boolean hasFilterCollector, boolean timeoutSet ) throws IOException { - // create the top docs collector last when the other collectors are known - final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, hasFilterCollector); - // add the top docs collector, the first collector context in the chain - collectors.addFirst(topDocsFactory); - final Collector queryCollector; if (searchContext.getProfilers() != null) { InternalProfileCollector profileCollector = QueryCollectorContext.createQueryCollectorWithProfiler(collectors); From 56911043ac4c7468f4ee0935e2b46b546e0b932d Mon Sep 17 00:00:00 2001 From: Luca Cavanna Date: Wed, 19 Apr 2023 18:42:10 +0200 Subject: [PATCH 2/7] spotless --- .../java/org/elasticsearch/search/query/QueryPhase.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java index 4db04bdfdd5f0..054f4c279ea58 100644 --- a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java @@ -128,8 +128,10 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut final LinkedList collectors = new LinkedList<>(); // create the top docs collector last when the other collectors are known - final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, - searchContext.parsedPostFilter() != null || searchContext.minimumScore() != null); + final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext( + searchContext, + searchContext.parsedPostFilter() != null || searchContext.minimumScore() != null + ); collectors.add(topDocsFactory); if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) { From 7f23f75cabc8db7187b88f65c550289c122b36b0 Mon Sep 17 00:00:00 2001 From: Luca Cavanna Date: Wed, 19 Apr 2023 23:53:33 +0200 Subject: [PATCH 3/7] Remove QueryCollectorContext abstraction --- .../search/profile/query/CollectorResult.java | 2 +- .../search/query/QueryCollectorContext.java | 170 ------------------ .../search/query/QueryPhase.java | 115 +++++++++--- .../search/query/TopDocsCollectorContext.java | 67 ++++--- .../search/query/QueryPhaseTests.java | 4 +- 5 files changed, 126 insertions(+), 232 deletions(-) delete mode 100644 server/src/main/java/org/elasticsearch/search/query/QueryCollectorContext.java diff --git a/server/src/main/java/org/elasticsearch/search/profile/query/CollectorResult.java b/server/src/main/java/org/elasticsearch/search/profile/query/CollectorResult.java index 6bc4e9706d5d0..4577bae6f24c1 100644 --- a/server/src/main/java/org/elasticsearch/search/profile/query/CollectorResult.java +++ b/server/src/main/java/org/elasticsearch/search/profile/query/CollectorResult.java @@ -67,7 +67,7 @@ public class CollectorResult implements ToXContentObject, Writeable { /** * A list of children collectors "embedded" inside this collector */ - private List children; + private final List children; public CollectorResult(String collectorName, String reason, long time, List children) { this.collectorName = collectorName; diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryCollectorContext.java b/server/src/main/java/org/elasticsearch/search/query/QueryCollectorContext.java deleted file mode 100644 index d53e035ac6dbf..0000000000000 --- a/server/src/main/java/org/elasticsearch/search/query/QueryCollectorContext.java +++ /dev/null @@ -1,170 +0,0 @@ -/* - * 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.search.query; - -import org.apache.lucene.search.Collector; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.MultiCollector; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.ScoreMode; -import org.apache.lucene.search.SimpleCollector; -import org.apache.lucene.search.Weight; -import org.elasticsearch.common.lucene.MinimumScoreCollector; -import org.elasticsearch.common.lucene.search.FilteredCollector; -import org.elasticsearch.search.profile.query.InternalProfileCollector; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; - -import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_MIN_SCORE; -import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_MULTI; -import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_POST_FILTER; -import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_TERMINATE_AFTER_COUNT; - -abstract class QueryCollectorContext { - private static final Collector EMPTY_COLLECTOR = new SimpleCollector() { - @Override - public void collect(int doc) {} - - @Override - public ScoreMode scoreMode() { - return ScoreMode.COMPLETE_NO_SCORES; - } - }; - - private final String profilerName; - - QueryCollectorContext(String profilerName) { - this.profilerName = profilerName; - } - - /** - * Creates a collector that delegates documents to the provided in collector. - * @param in The delegate collector - */ - abstract Collector create(Collector in) throws IOException; - - /** - * Wraps this collector with a profiler - */ - protected InternalProfileCollector createWithProfiler(InternalProfileCollector in) throws IOException { - final Collector collector = create(in); - return new InternalProfileCollector(collector, profilerName, in != null ? Collections.singletonList(in) : Collections.emptyList()); - } - - /** - * Post-process result after search execution. - * - * @param result The query search result to populate - */ - void postProcess(QuerySearchResult result) throws IOException {} - - /** - * Creates the collector tree from the provided collectors - * @param collectors Ordered list of collector context - */ - static Collector createQueryCollector(List collectors) throws IOException { - Collector collector = null; - for (QueryCollectorContext ctx : collectors) { - collector = ctx.create(collector); - } - return collector; - } - - /** - * Creates the collector tree from the provided collectors and wraps each collector with a profiler - * @param collectors Ordered list of collector context - */ - static InternalProfileCollector createQueryCollectorWithProfiler(List collectors) throws IOException { - InternalProfileCollector collector = null; - for (QueryCollectorContext ctx : collectors) { - collector = ctx.createWithProfiler(collector); - } - return collector; - } - - /** - * Filters documents with a query score greater than minScore - * @param minScore The minimum score filter - */ - static QueryCollectorContext createMinScoreCollectorContext(float minScore) { - return new QueryCollectorContext(REASON_SEARCH_MIN_SCORE) { - @Override - Collector create(Collector in) { - return new MinimumScoreCollector(in, minScore); - } - }; - } - - /** - * Filters documents based on the provided query - */ - static QueryCollectorContext createFilteredCollectorContext(IndexSearcher searcher, Query query) { - return new QueryCollectorContext(REASON_SEARCH_POST_FILTER) { - @Override - Collector create(Collector in) throws IOException { - final Weight filterWeight = searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1f); - return new FilteredCollector(in, filterWeight); - } - }; - } - - /** - * Creates a multi collector from the provided sub-collector - */ - static QueryCollectorContext createAggsCollectorContext(Collector subCollector) { - return new QueryCollectorContext(REASON_SEARCH_MULTI) { - @Override - Collector create(Collector in) { - List subCollectors = new ArrayList<>(); - subCollectors.add(in); - subCollectors.add(subCollector); - return MultiCollector.wrap(subCollectors); - } - - @Override - protected InternalProfileCollector createWithProfiler(InternalProfileCollector in) { - final List subCollectors = new ArrayList<>(); - subCollectors.add(in); - if (subCollector instanceof InternalProfileCollector == false) { - throw new IllegalArgumentException("non-profiling collector"); - } - subCollectors.add((InternalProfileCollector) subCollector); - final Collector collector = MultiCollector.wrap(subCollectors); - return new InternalProfileCollector(collector, REASON_SEARCH_MULTI, subCollectors); - } - }; - } - - /** - * Creates collector limiting the collection to the first numHits documents - */ - static QueryCollectorContext createEarlyTerminationCollectorContext(int numHits) { - return new QueryCollectorContext(REASON_SEARCH_TERMINATE_AFTER_COUNT) { - private Collector collector; - - /** - * Creates a {@link MultiCollector} to ensure that the {@link EarlyTerminatingCollector} - * can terminate the collection independently of the provided in {@link Collector}. - */ - @Override - Collector create(Collector in) { - assert collector == null; - - List subCollectors = new ArrayList<>(); - subCollectors.add(new EarlyTerminatingCollector(EMPTY_COLLECTOR, numHits, true)); - subCollectors.add(in); - this.collector = MultiCollector.wrap(subCollectors); - return collector; - } - }; - } -} diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java index 054f4c279ea58..de95ab506cb27 100644 --- a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java @@ -16,12 +16,18 @@ import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.search.Collector; import org.apache.lucene.search.FieldDoc; +import org.apache.lucene.search.MultiCollector; import org.apache.lucene.search.Query; import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.SimpleCollector; import org.apache.lucene.search.Sort; import org.apache.lucene.search.TopDocs; import org.apache.lucene.search.TotalHits; +import org.apache.lucene.search.Weight; import org.elasticsearch.common.lucene.Lucene; +import org.elasticsearch.common.lucene.MinimumScoreCollector; +import org.elasticsearch.common.lucene.search.FilteredCollector; import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore; import org.elasticsearch.common.util.concurrent.EWMATrackingEsThreadPoolExecutor; import org.elasticsearch.common.util.concurrent.EsThreadPoolExecutor; @@ -33,6 +39,7 @@ import org.elasticsearch.search.internal.ContextIndexSearcher; import org.elasticsearch.search.internal.ScrollContext; import org.elasticsearch.search.internal.SearchContext; +import org.elasticsearch.search.profile.Profilers; import org.elasticsearch.search.profile.query.InternalProfileCollector; import org.elasticsearch.search.rescore.RescorePhase; import org.elasticsearch.search.sort.SortAndFormats; @@ -40,13 +47,15 @@ import org.elasticsearch.threadpool.ThreadPool; import java.io.IOException; -import java.util.LinkedList; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.concurrent.ExecutorService; -import static org.elasticsearch.search.query.QueryCollectorContext.createAggsCollectorContext; -import static org.elasticsearch.search.query.QueryCollectorContext.createEarlyTerminationCollectorContext; -import static org.elasticsearch.search.query.QueryCollectorContext.createFilteredCollectorContext; -import static org.elasticsearch.search.query.QueryCollectorContext.createMinScoreCollectorContext; +import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_MIN_SCORE; +import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_MULTI; +import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_POST_FILTER; +import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_TERMINATE_AFTER_COUNT; import static org.elasticsearch.search.query.TopDocsCollectorContext.createTopDocsCollectorContext; /** @@ -125,32 +134,71 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut } } - final LinkedList collectors = new LinkedList<>(); - // create the top docs collector last when the other collectors are known final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext( searchContext, searchContext.parsedPostFilter() != null || searchContext.minimumScore() != null ); - collectors.add(topDocsFactory); + + Collector collector = wrapWithProfilerCollectorIfNeeded( + searchContext.getProfilers(), + topDocsFactory.collector(), + topDocsFactory.profilerName, + null + ); + ; if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) { // add terminate_after before the filter collectors // it will only be applied on documents accepted by these filter collectors - collectors.add(createEarlyTerminationCollectorContext(searchContext.terminateAfter())); + List subCollectors = new ArrayList<>(); + subCollectors.add(new EarlyTerminatingCollector(EMPTY_COLLECTOR, searchContext.terminateAfter(), true)); + subCollectors.add(collector); + collector = wrapWithProfilerCollectorIfNeeded( + searchContext.getProfilers(), + MultiCollector.wrap(subCollectors), + REASON_SEARCH_TERMINATE_AFTER_COUNT, + collector + ); } if (searchContext.parsedPostFilter() != null) { // add post filters before aggregations // it will only be applied to top hits - collectors.add(createFilteredCollectorContext(searcher, searchContext.parsedPostFilter().query())); + final Weight filterWeight = searcher.createWeight( + searcher.rewrite(searchContext.parsedPostFilter().query()), + ScoreMode.COMPLETE_NO_SCORES, + 1f + ); + collector = wrapWithProfilerCollectorIfNeeded( + searchContext.getProfilers(), + new FilteredCollector(collector, filterWeight), + REASON_SEARCH_POST_FILTER, + collector + ); } if (searchContext.getAggsCollector() != null) { // plug in additional collectors, like aggregations - collectors.add(createAggsCollectorContext(searchContext.getAggsCollector())); + List subCollectors = new ArrayList<>(); + subCollectors.add(collector); + subCollectors.add(searchContext.getAggsCollector()); + collector = MultiCollector.wrap(subCollectors); + if (searchContext.getProfilers() != null) { + List profileCollectors = List.of( + (InternalProfileCollector) collector, + (InternalProfileCollector) searchContext.getAggsCollector() + ); + // in this case we pass both collectors as children profile collectors + collector = new InternalProfileCollector(collector, REASON_SEARCH_MULTI, profileCollectors); + } } if (searchContext.minimumScore() != null) { // apply the minimum score after multi collector so we filter aggs as well - collectors.add(createMinScoreCollectorContext(searchContext.minimumScore())); + collector = wrapWithProfilerCollectorIfNeeded( + searchContext.getProfilers(), + new MinimumScoreCollector(collector, searchContext.minimumScore()), + REASON_SEARCH_MIN_SCORE, + collector + ); } boolean timeoutSet = scrollContext == null @@ -173,7 +221,8 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut } try { - searchWithCollector(searchContext, searcher, query, collectors, timeoutSet); + searchWithCollector(searchContext, searcher, query, collector, timeoutSet); + queryResult.topDocs(topDocsFactory.topDocsAndMaxScore(), topDocsFactory.sortValueFormats); ExecutorService executor = searchContext.indexShard().getThreadPool().executor(ThreadPool.Names.SEARCH); assert executor instanceof EWMATrackingEsThreadPoolExecutor || (executor instanceof EsThreadPoolExecutor == false /* in case thread pool is mocked out in tests */) @@ -194,24 +243,35 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut } } + private static Collector wrapWithProfilerCollectorIfNeeded( + Profilers profilers, + Collector collector, + String profilerName, + Collector child + ) { + if (profilers == null) { + return collector; + } + return new InternalProfileCollector( + collector, + profilerName, + child != null ? Collections.singletonList((InternalProfileCollector) child) : Collections.emptyList() + ); + } + private static void searchWithCollector( SearchContext searchContext, ContextIndexSearcher searcher, Query query, - LinkedList collectors, + Collector collector, boolean timeoutSet ) throws IOException { - final Collector queryCollector; if (searchContext.getProfilers() != null) { - InternalProfileCollector profileCollector = QueryCollectorContext.createQueryCollectorWithProfiler(collectors); - searchContext.getProfilers().getCurrentQueryProfiler().setCollector(profileCollector); - queryCollector = profileCollector; - } else { - queryCollector = QueryCollectorContext.createQueryCollector(collectors); + searchContext.getProfilers().getCurrentQueryProfiler().setCollector((InternalProfileCollector) collector); } QuerySearchResult queryResult = searchContext.queryResult(); try { - searcher.search(query, queryCollector); + searcher.search(query, collector); } catch (EarlyTerminatingCollector.EarlyTerminationException e) { queryResult.terminatedEarly(true); } catch (TimeExceededException e) { @@ -225,9 +285,6 @@ private static void searchWithCollector( if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER && queryResult.terminatedEarly() == null) { queryResult.terminatedEarly(false); } - for (QueryCollectorContext ctx : collectors) { - ctx.postProcess(queryResult); - } } /** @@ -249,4 +306,14 @@ private static boolean canEarlyTerminate(IndexReader reader, SortAndFormats sort } public static class TimeExceededException extends RuntimeException {} + + private static final Collector EMPTY_COLLECTOR = new SimpleCollector() { + @Override + public void collect(int doc) {} + + @Override + public ScoreMode scoreMode() { + return ScoreMode.COMPLETE_NO_SCORES; + } + }; } diff --git a/server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java b/server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java index 7bce4d83b10a8..7fb2991938837 100644 --- a/server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java +++ b/server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java @@ -64,16 +64,21 @@ import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_TOP_HITS; /** - * A {@link QueryCollectorContext} that creates top docs collector + * Creates and holds the main {@link Collector} that will be used to search and collect top hits. */ -abstract class TopDocsCollectorContext extends QueryCollectorContext { - protected final int numHits; +abstract class TopDocsCollectorContext { + final String profilerName; + final DocValueFormat[] sortValueFormats; - TopDocsCollectorContext(String profilerName, int numHits) { - super(profilerName); - this.numHits = numHits; + TopDocsCollectorContext(String profilerName, DocValueFormat[] sortValueFormats) { + this.profilerName = profilerName; + this.sortValueFormats = sortValueFormats; } + abstract Collector collector(); + + abstract TopDocsAndMaxScore topDocsAndMaxScore() throws IOException; + static class EmptyTopDocsCollectorContext extends TopDocsCollectorContext { private final Sort sort; private final Collector collector; @@ -85,7 +90,7 @@ static class EmptyTopDocsCollectorContext extends TopDocsCollectorContext { * @param trackTotalHitsUpTo The threshold up to which total hit count needs to be tracked */ private EmptyTopDocsCollectorContext(@Nullable SortAndFormats sortAndFormats, int trackTotalHitsUpTo) { - super(REASON_SEARCH_COUNT, 0); + super(REASON_SEARCH_COUNT, null); this.sort = sortAndFormats == null ? null : sortAndFormats.sort; if (trackTotalHitsUpTo == SearchContext.TRACK_TOTAL_HITS_DISABLED) { this.collector = new EarlyTerminatingCollector(new TotalHitCountCollector(), 0, false); @@ -108,13 +113,12 @@ private EmptyTopDocsCollectorContext(@Nullable SortAndFormats sortAndFormats, in } @Override - Collector create(Collector in) { - assert in == null; + Collector collector() { return collector; } @Override - void postProcess(QuerySearchResult result) { + TopDocsAndMaxScore topDocsAndMaxScore() { final TotalHits totalHitCount = hitCountSupplier.get(); final TopDocs topDocs; if (sort != null) { @@ -122,12 +126,11 @@ void postProcess(QuerySearchResult result) { } else { topDocs = new TopDocs(totalHitCount, Lucene.EMPTY_SCORE_DOCS); } - result.topDocs(new TopDocsAndMaxScore(topDocs, Float.NaN), null); + return new TopDocsAndMaxScore(topDocs, Float.NaN); } } static class CollapsingTopDocsCollectorContext extends TopDocsCollectorContext { - private final DocValueFormat[] sortFmt; private final SinglePassGroupingCollector topDocsCollector; private final Supplier maxScoreSupplier; @@ -145,11 +148,10 @@ private CollapsingTopDocsCollectorContext( boolean trackMaxScore, @Nullable FieldDoc after ) { - super(REASON_SEARCH_TOP_HITS, numHits); + super(REASON_SEARCH_TOP_HITS, sortAndFormats == null ? new DocValueFormat[] { DocValueFormat.RAW } : sortAndFormats.formats); assert numHits > 0; assert collapseContext != null; Sort sort = sortAndFormats == null ? Sort.RELEVANCE : sortAndFormats.sort; - this.sortFmt = sortAndFormats == null ? new DocValueFormat[] { DocValueFormat.RAW } : sortAndFormats.formats; this.topDocsCollector = collapseContext.createTopDocs(sort, numHits, after); MaxScoreCollector maxScoreCollector; @@ -162,15 +164,14 @@ private CollapsingTopDocsCollectorContext( } @Override - Collector create(Collector in) throws IOException { - assert in == null; + Collector collector() { return topDocsCollector; } @Override - void postProcess(QuerySearchResult result) throws IOException { + TopDocsAndMaxScore topDocsAndMaxScore() throws IOException { TopFieldGroups topDocs = topDocsCollector.getTopGroups(0); - result.topDocs(new TopDocsAndMaxScore(topDocs, maxScoreSupplier.get()), sortFmt); + return new TopDocsAndMaxScore(topDocs, maxScoreSupplier.get()); } } @@ -189,7 +190,6 @@ private static TopDocsCollector createCollector( } } - protected final @Nullable SortAndFormats sortAndFormats; private final Collector collector; private final Supplier totalHitsSupplier; private final Supplier topDocsSupplier; @@ -197,12 +197,13 @@ private static TopDocsCollector createCollector( /** * Ctr - * @param reader The index reader - * @param query The Lucene query - * @param sortAndFormats The sort clause if provided - * @param numHits The number of top hits to retrieve - * @param searchAfter The doc this request should "search after" - * @param trackMaxScore True if max score should be tracked + * + * @param reader The index reader + * @param query The Lucene query + * @param sortAndFormats The sort clause if provided + * @param numHits The number of top hits to retrieve + * @param searchAfter The doc this request should "search after" + * @param trackMaxScore True if max score should be tracked * @param trackTotalHitsUpTo Threshold up to which total hit count should be tracked * @param hasFilterCollector True if the collector chain contains at least one collector that can filter documents out */ @@ -216,11 +217,9 @@ private SimpleTopDocsCollectorContext( int trackTotalHitsUpTo, boolean hasFilterCollector ) throws IOException { - super(REASON_SEARCH_TOP_HITS, numHits); - this.sortAndFormats = sortAndFormats; + super(REASON_SEARCH_TOP_HITS, sortAndFormats == null ? null : sortAndFormats.formats); final TopDocsCollector topDocsCollector; - if ((sortAndFormats == null || SortField.FIELD_SCORE.equals(sortAndFormats.sort.getSort()[0])) && hasInfMaxScore(query)) { // disable max score optimization since we have a mandatory clause // that doesn't track the maximum score @@ -268,8 +267,7 @@ private SimpleTopDocsCollectorContext( } @Override - Collector create(Collector in) { - assert in == null; + Collector collector() { return collector; } @@ -286,9 +284,8 @@ TopDocsAndMaxScore newTopDocs() { } @Override - void postProcess(QuerySearchResult result) throws IOException { - final TopDocsAndMaxScore topDocs = newTopDocs(); - result.topDocs(topDocs, sortAndFormats == null ? null : sortAndFormats.formats); + TopDocsAndMaxScore topDocsAndMaxScore() { + return newTopDocs(); } } @@ -322,7 +319,7 @@ private ScrollingTopDocsCollectorContext( } @Override - void postProcess(QuerySearchResult result) throws IOException { + TopDocsAndMaxScore topDocsAndMaxScore() { final TopDocsAndMaxScore topDocs = newTopDocs(); if (scrollContext.totalHits == null) { // first round @@ -341,7 +338,7 @@ void postProcess(QuerySearchResult result) throws IOException { scrollContext.lastEmittedDoc = topDocs.topDocs.scoreDocs[topDocs.topDocs.scoreDocs.length - 1]; } } - result.topDocs(topDocs, sortAndFormats == null ? null : sortAndFormats.formats); + return topDocs; } } diff --git a/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java index c3e77088b10fd..943bec8c8d1cc 100644 --- a/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java +++ b/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java @@ -680,7 +680,7 @@ public void testDisableTopScoreCollection() throws Exception { context.setSize(3); context.trackTotalHitsUpTo(3); TopDocsCollectorContext topDocsContext = TopDocsCollectorContext.createTopDocsCollectorContext(context, false); - assertEquals(topDocsContext.create(null).scoreMode(), org.apache.lucene.search.ScoreMode.COMPLETE); + assertEquals(topDocsContext.collector().scoreMode(), org.apache.lucene.search.ScoreMode.COMPLETE); QueryPhase.executeInternal(context); assertEquals(5, context.queryResult().topDocs().topDocs.totalHits.value); assertEquals(context.queryResult().topDocs().topDocs.totalHits.relation, TotalHits.Relation.EQUAL_TO); @@ -688,7 +688,7 @@ public void testDisableTopScoreCollection() throws Exception { context.sort(new SortAndFormats(new Sort(new SortField("other", SortField.Type.INT)), new DocValueFormat[] { DocValueFormat.RAW })); topDocsContext = TopDocsCollectorContext.createTopDocsCollectorContext(context, false); - assertEquals(topDocsContext.create(null).scoreMode(), org.apache.lucene.search.ScoreMode.TOP_DOCS); + assertEquals(topDocsContext.collector().scoreMode(), org.apache.lucene.search.ScoreMode.TOP_DOCS); QueryPhase.executeInternal(context); assertEquals(5, context.queryResult().topDocs().topDocs.totalHits.value); assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(3)); From c5c3bcd12ad2b1cf612baa274281827e7eb151c2 Mon Sep 17 00:00:00 2001 From: Luca Cavanna Date: Thu, 20 Apr 2023 00:09:16 +0200 Subject: [PATCH 4/7] Fix class cast --- .../java/org/elasticsearch/search/query/QueryPhase.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java index de95ab506cb27..21cfe6b580624 100644 --- a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java @@ -146,7 +146,6 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut topDocsFactory.profilerName, null ); - ; if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) { // add terminate_after before the filter collectors @@ -181,14 +180,15 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut List subCollectors = new ArrayList<>(); subCollectors.add(collector); subCollectors.add(searchContext.getAggsCollector()); - collector = MultiCollector.wrap(subCollectors); - if (searchContext.getProfilers() != null) { + if (searchContext.getProfilers() == null) { + collector = MultiCollector.wrap(subCollectors); + } else { List profileCollectors = List.of( (InternalProfileCollector) collector, (InternalProfileCollector) searchContext.getAggsCollector() ); // in this case we pass both collectors as children profile collectors - collector = new InternalProfileCollector(collector, REASON_SEARCH_MULTI, profileCollectors); + collector = new InternalProfileCollector(MultiCollector.wrap(subCollectors), REASON_SEARCH_MULTI, profileCollectors); } } if (searchContext.minimumScore() != null) { From 1513757d20cb375154f49cf98ecf27ea9de6344c Mon Sep 17 00:00:00 2001 From: Luca Cavanna Date: Thu, 20 Apr 2023 20:32:43 +0200 Subject: [PATCH 5/7] iter --- .../java/org/elasticsearch/search/query/QueryPhase.java | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java index 21cfe6b580624..7e687b0aff795 100644 --- a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java @@ -176,19 +176,16 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut ); } if (searchContext.getAggsCollector() != null) { - // plug in additional collectors, like aggregations - List subCollectors = new ArrayList<>(); - subCollectors.add(collector); - subCollectors.add(searchContext.getAggsCollector()); if (searchContext.getProfilers() == null) { - collector = MultiCollector.wrap(subCollectors); + collector = MultiCollector.wrap(collector, searchContext.getAggsCollector()); } else { List profileCollectors = List.of( (InternalProfileCollector) collector, (InternalProfileCollector) searchContext.getAggsCollector() ); // in this case we pass both collectors as children profile collectors - collector = new InternalProfileCollector(MultiCollector.wrap(subCollectors), REASON_SEARCH_MULTI, profileCollectors); + collector = new InternalProfileCollector(MultiCollector.wrap(collector, searchContext.getAggsCollector()), + REASON_SEARCH_MULTI, profileCollectors); } } if (searchContext.minimumScore() != null) { From 9f7d9d0cbb3d27ea46ebda27f2bfa1dc182d7cdd Mon Sep 17 00:00:00 2001 From: Luca Cavanna Date: Thu, 20 Apr 2023 23:12:41 +0200 Subject: [PATCH 6/7] iter --- .../query/InternalProfileCollector.java | 8 +++--- .../search/query/QueryPhase.java | 28 ++++++------------- 2 files changed, 13 insertions(+), 23 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollector.java b/server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollector.java index 5eb5c057c1479..c4914da8fee0f 100644 --- a/server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollector.java +++ b/server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollector.java @@ -47,15 +47,15 @@ public class InternalProfileCollector implements Collector { */ private final InternalProfileCollector[] children; - public InternalProfileCollector(Collector collector, String reason, InternalProfileCollector... children) { + public InternalProfileCollector(Collector collector, String reason, Collector... children) { this.collector = new ProfileCollector(collector); this.reason = reason; this.collectorName = deriveCollectorName(collector); Objects.requireNonNull(children, "children collectors cannot be null"); - for (InternalProfileCollector child : children) { - Objects.requireNonNull(child, "child collector cannot be null"); + this.children = new InternalProfileCollector[children.length]; + for (int i = 0; i < children.length; i++) { + this.children[i] = (InternalProfileCollector) Objects.requireNonNull(children[i], "child collector cannot be null"); } - this.children = children; } /** diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java index 7e687b0aff795..354dab02f3e76 100644 --- a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java @@ -47,8 +47,7 @@ import org.elasticsearch.threadpool.ThreadPool; import java.io.IOException; -import java.util.ArrayList; -import java.util.Collections; +import java.util.Arrays; import java.util.List; import java.util.concurrent.ExecutorService; @@ -150,12 +149,11 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) { // add terminate_after before the filter collectors // it will only be applied on documents accepted by these filter collectors - List subCollectors = new ArrayList<>(); - subCollectors.add(new EarlyTerminatingCollector(EMPTY_COLLECTOR, searchContext.terminateAfter(), true)); - subCollectors.add(collector); + EarlyTerminatingCollector earlyTerminatingCollector = new EarlyTerminatingCollector(EMPTY_COLLECTOR, + searchContext.terminateAfter(), true); collector = wrapWithProfilerCollectorIfNeeded( searchContext.getProfilers(), - MultiCollector.wrap(subCollectors), + MultiCollector.wrap(earlyTerminatingCollector, collector), REASON_SEARCH_TERMINATE_AFTER_COUNT, collector ); @@ -176,17 +174,9 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut ); } if (searchContext.getAggsCollector() != null) { - if (searchContext.getProfilers() == null) { - collector = MultiCollector.wrap(collector, searchContext.getAggsCollector()); - } else { - List profileCollectors = List.of( - (InternalProfileCollector) collector, - (InternalProfileCollector) searchContext.getAggsCollector() - ); - // in this case we pass both collectors as children profile collectors - collector = new InternalProfileCollector(MultiCollector.wrap(collector, searchContext.getAggsCollector()), - REASON_SEARCH_MULTI, profileCollectors); - } + collector = wrapWithProfilerCollectorIfNeeded(searchContext.getProfilers(), + MultiCollector.wrap(collector, searchContext.getAggsCollector()), + REASON_SEARCH_MULTI, collector, searchContext.getAggsCollector()); } if (searchContext.minimumScore() != null) { // apply the minimum score after multi collector so we filter aggs as well @@ -244,7 +234,7 @@ private static Collector wrapWithProfilerCollectorIfNeeded( Profilers profilers, Collector collector, String profilerName, - Collector child + Collector... children ) { if (profilers == null) { return collector; @@ -252,7 +242,7 @@ private static Collector wrapWithProfilerCollectorIfNeeded( return new InternalProfileCollector( collector, profilerName, - child != null ? Collections.singletonList((InternalProfileCollector) child) : Collections.emptyList() + children ); } From 1e69545da2fe48ab0556f118d0473b27b13f6dff Mon Sep 17 00:00:00 2001 From: Luca Cavanna Date: Thu, 20 Apr 2023 23:19:28 +0200 Subject: [PATCH 7/7] spotless --- .../search/query/QueryPhase.java | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java index 354dab02f3e76..49ba83b9466f1 100644 --- a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java @@ -47,8 +47,6 @@ import org.elasticsearch.threadpool.ThreadPool; import java.io.IOException; -import java.util.Arrays; -import java.util.List; import java.util.concurrent.ExecutorService; import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_MIN_SCORE; @@ -142,15 +140,17 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut Collector collector = wrapWithProfilerCollectorIfNeeded( searchContext.getProfilers(), topDocsFactory.collector(), - topDocsFactory.profilerName, - null + topDocsFactory.profilerName ); if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) { // add terminate_after before the filter collectors // it will only be applied on documents accepted by these filter collectors - EarlyTerminatingCollector earlyTerminatingCollector = new EarlyTerminatingCollector(EMPTY_COLLECTOR, - searchContext.terminateAfter(), true); + EarlyTerminatingCollector earlyTerminatingCollector = new EarlyTerminatingCollector( + EMPTY_COLLECTOR, + searchContext.terminateAfter(), + true + ); collector = wrapWithProfilerCollectorIfNeeded( searchContext.getProfilers(), MultiCollector.wrap(earlyTerminatingCollector, collector), @@ -174,9 +174,13 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut ); } if (searchContext.getAggsCollector() != null) { - collector = wrapWithProfilerCollectorIfNeeded(searchContext.getProfilers(), + collector = wrapWithProfilerCollectorIfNeeded( + searchContext.getProfilers(), MultiCollector.wrap(collector, searchContext.getAggsCollector()), - REASON_SEARCH_MULTI, collector, searchContext.getAggsCollector()); + REASON_SEARCH_MULTI, + collector, + searchContext.getAggsCollector() + ); } if (searchContext.minimumScore() != null) { // apply the minimum score after multi collector so we filter aggs as well @@ -239,11 +243,7 @@ private static Collector wrapWithProfilerCollectorIfNeeded( if (profilers == null) { return collector; } - return new InternalProfileCollector( - collector, - profilerName, - children - ); + return new InternalProfileCollector(collector, profilerName, children); } private static void searchWithCollector(