diff --git a/CHANGELOG.md b/CHANGELOG.md index 67ce05d4c1bc9..f8760bb96ee68 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -129,6 +129,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Changed `opensearch-env` to respect already set `OPENSEARCH_HOME` environment variable ([#6956](https://github.com/opensearch-project/OpenSearch/pull/6956/)) - Added a new field type: flat_object ([#6507](https://github.com/opensearch-project/OpenSearch/pull/6507)) - Increased visibility of BaseRestHandler’s `unrecognized` method using a new public `unrecognizedStrings` method. ([#7125](https://github.com/opensearch-project/OpenSearch/pull/7125)) +- Moved concurrent-search from sandbox plugin to server module behind feature flag ([#7203](https://github.com/opensearch-project/OpenSearch/pull/7203)) ### Deprecated - Map, List, and Set in org.opensearch.common.collect ([#6609](https://github.com/opensearch-project/OpenSearch/pull/6609)) diff --git a/sandbox/plugins/concurrent-search/build.gradle b/sandbox/plugins/concurrent-search/build.gradle deleted file mode 100644 index 0e766dc4fc1ba..0000000000000 --- a/sandbox/plugins/concurrent-search/build.gradle +++ /dev/null @@ -1,42 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - * - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. - */ - -/* - * Licensed to Elasticsearch under one or more contributor - * license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright - * ownership. Elasticsearch licenses this file to you under - * the Apache License, Version 2.0 (the "License"); you may - * not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -apply plugin: 'opensearch.opensearchplugin' -apply plugin: 'opensearch.yaml-rest-test' - -opensearchplugin { - name 'concurrent-search' - description 'The experimental plugin which implements concurrent search over Apache Lucene segments' - classname 'org.opensearch.search.ConcurrentSegmentSearchPlugin' - licenseFile rootProject.file('licenses/APACHE-LICENSE-2.0.txt') - noticeFile rootProject.file('NOTICE.txt') -} - -yamlRestTest.enabled = false; -testingConventions.enabled = false; diff --git a/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/ConcurrentSegmentSearchPlugin.java b/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/ConcurrentSegmentSearchPlugin.java deleted file mode 100644 index da999e40f0f07..0000000000000 --- a/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/ConcurrentSegmentSearchPlugin.java +++ /dev/null @@ -1,53 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.search; - -import org.opensearch.common.settings.Settings; -import org.opensearch.common.util.concurrent.OpenSearchExecutors; -import org.opensearch.plugins.Plugin; -import org.opensearch.plugins.SearchPlugin; -import org.opensearch.search.query.ConcurrentQueryPhaseSearcher; -import org.opensearch.search.query.QueryPhaseSearcher; -import org.opensearch.threadpool.ExecutorBuilder; -import org.opensearch.threadpool.FixedExecutorBuilder; -import org.opensearch.threadpool.ThreadPool; - -import java.util.Collections; -import java.util.List; -import java.util.Optional; - -/** - * The experimental plugin which implements the concurrent search over Apache Lucene segments. - */ -public class ConcurrentSegmentSearchPlugin extends Plugin implements SearchPlugin { - private static final String INDEX_SEARCHER = "index_searcher"; - - /** - * Default constructor - */ - public ConcurrentSegmentSearchPlugin() {} - - @Override - public Optional getQueryPhaseSearcher() { - return Optional.of(new ConcurrentQueryPhaseSearcher()); - } - - @Override - public List> getExecutorBuilders(Settings settings) { - final int allocatedProcessors = OpenSearchExecutors.allocatedProcessors(settings); - return Collections.singletonList( - new FixedExecutorBuilder(settings, INDEX_SEARCHER, allocatedProcessors, 1000, "thread_pool." + INDEX_SEARCHER) - ); - } - - @Override - public Optional getIndexSearcherExecutorProvider() { - return Optional.of((ThreadPool threadPool) -> threadPool.executor(INDEX_SEARCHER)); - } -} diff --git a/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/package-info.java b/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/package-info.java deleted file mode 100644 index 041f914fab7d7..0000000000000 --- a/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/package-info.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -/** - * The implementation of the experimental plugin which implements the concurrent search over Apache Lucene segments. - */ -package org.opensearch.search; diff --git a/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/query/package-info.java b/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/query/package-info.java deleted file mode 100644 index 0f98ae7682a84..0000000000000 --- a/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/query/package-info.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -/** - * {@link org.opensearch.search.query.QueryPhaseSearcher} implementation for concurrent search - */ -package org.opensearch.search.query; diff --git a/sandbox/plugins/concurrent-search/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java b/sandbox/plugins/concurrent-search/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java deleted file mode 100644 index 0b88abb314fcd..0000000000000 --- a/sandbox/plugins/concurrent-search/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java +++ /dev/null @@ -1,316 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.search.profile.query; - -import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; - -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field.Store; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.IndexReader; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.Term; -import org.apache.lucene.search.Explanation; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.LRUQueryCache; -import org.apache.lucene.search.LeafCollector; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.QueryCachingPolicy; -import org.apache.lucene.search.QueryVisitor; -import org.apache.lucene.search.ScoreMode; -import org.apache.lucene.search.Scorer; -import org.apache.lucene.search.ScorerSupplier; -import org.apache.lucene.search.Sort; -import org.apache.lucene.search.TermQuery; -import org.apache.lucene.search.TotalHitCountCollector; -import org.apache.lucene.search.Weight; -import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.index.RandomIndexWriter; -import org.apache.lucene.tests.search.RandomApproximationQuery; -import org.apache.lucene.tests.util.TestUtil; -import org.opensearch.common.util.io.IOUtils; -import org.opensearch.search.internal.ContextIndexSearcher; -import org.opensearch.search.profile.ProfileResult; -import org.opensearch.test.OpenSearchTestCase; -import org.opensearch.threadpool.ThreadPool; -import org.junit.After; -import org.junit.Before; - -import java.io.IOException; -import java.util.Arrays; -import java.util.Collection; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; - -import static org.hamcrest.Matchers.equalTo; -import static org.hamcrest.Matchers.greaterThan; - -public class QueryProfilerTests extends OpenSearchTestCase { - - private Directory dir; - private IndexReader reader; - private ContextIndexSearcher searcher; - private ExecutorService executor; - - @ParametersFactory - public static Collection concurrency() { - return Arrays.asList(new Integer[] { 0 }, new Integer[] { 5 }); - } - - public QueryProfilerTests(int concurrency) { - this.executor = (concurrency > 0) ? Executors.newFixedThreadPool(concurrency) : null; - } - - @Before - public void setUp() throws Exception { - super.setUp(); - - dir = newDirectory(); - RandomIndexWriter w = new RandomIndexWriter(random(), dir); - final int numDocs = TestUtil.nextInt(random(), 1, 20); - for (int i = 0; i < numDocs; ++i) { - final int numHoles = random().nextInt(5); - for (int j = 0; j < numHoles; ++j) { - w.addDocument(new Document()); - } - Document doc = new Document(); - doc.add(new StringField("foo", "bar", Store.NO)); - w.addDocument(doc); - } - reader = w.getReader(); - w.close(); - searcher = new ContextIndexSearcher( - reader, - IndexSearcher.getDefaultSimilarity(), - IndexSearcher.getDefaultQueryCache(), - ALWAYS_CACHE_POLICY, - true, - executor - ); - } - - @After - public void tearDown() throws Exception { - super.tearDown(); - - LRUQueryCache cache = (LRUQueryCache) searcher.getQueryCache(); - assertThat(cache.getHitCount(), equalTo(0L)); - assertThat(cache.getCacheCount(), equalTo(0L)); - assertThat(cache.getTotalCount(), equalTo(cache.getMissCount())); - assertThat(cache.getCacheSize(), equalTo(0L)); - - if (executor != null) { - ThreadPool.terminate(executor, 10, TimeUnit.SECONDS); - } - - IOUtils.close(reader, dir); - dir = null; - reader = null; - searcher = null; - } - - public void testBasic() throws IOException { - QueryProfiler profiler = new QueryProfiler(executor != null); - searcher.setProfiler(profiler); - Query query = new TermQuery(new Term("foo", "bar")); - searcher.search(query, 1); - List results = profiler.getTree(); - assertEquals(1, results.size()); - Map breakdown = results.get(0).getTimeBreakdown(); - assertThat(breakdown.get(QueryTimingType.CREATE_WEIGHT.toString()), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.BUILD_SCORER.toString()), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.NEXT_DOC.toString()), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.ADVANCE.toString()), equalTo(0L)); - assertThat(breakdown.get(QueryTimingType.SCORE.toString()), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.MATCH.toString()), equalTo(0L)); - - assertThat(breakdown.get(QueryTimingType.CREATE_WEIGHT.toString() + "_count"), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.BUILD_SCORER.toString() + "_count"), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.NEXT_DOC.toString() + "_count"), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.ADVANCE.toString() + "_count"), equalTo(0L)); - assertThat(breakdown.get(QueryTimingType.SCORE.toString() + "_count"), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.MATCH.toString() + "_count"), equalTo(0L)); - - long rewriteTime = profiler.getRewriteTime(); - assertThat(rewriteTime, greaterThan(0L)); - } - - public void testNoScoring() throws IOException { - QueryProfiler profiler = new QueryProfiler(executor != null); - searcher.setProfiler(profiler); - Query query = new TermQuery(new Term("foo", "bar")); - searcher.search(query, 1, Sort.INDEXORDER); // scores are not needed - List results = profiler.getTree(); - assertEquals(1, results.size()); - Map breakdown = results.get(0).getTimeBreakdown(); - assertThat(breakdown.get(QueryTimingType.CREATE_WEIGHT.toString()), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.BUILD_SCORER.toString()), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.NEXT_DOC.toString()), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.ADVANCE.toString()), equalTo(0L)); - assertThat(breakdown.get(QueryTimingType.SCORE.toString()), equalTo(0L)); - assertThat(breakdown.get(QueryTimingType.MATCH.toString()), equalTo(0L)); - - assertThat(breakdown.get(QueryTimingType.CREATE_WEIGHT.toString() + "_count"), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.BUILD_SCORER.toString() + "_count"), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.NEXT_DOC.toString() + "_count"), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.ADVANCE.toString() + "_count"), equalTo(0L)); - assertThat(breakdown.get(QueryTimingType.SCORE.toString() + "_count"), equalTo(0L)); - assertThat(breakdown.get(QueryTimingType.MATCH.toString() + "_count"), equalTo(0L)); - - long rewriteTime = profiler.getRewriteTime(); - assertThat(rewriteTime, greaterThan(0L)); - } - - public void testUseIndexStats() throws IOException { - QueryProfiler profiler = new QueryProfiler(executor != null); - searcher.setProfiler(profiler); - Query query = new TermQuery(new Term("foo", "bar")); - searcher.count(query); // will use index stats - List results = profiler.getTree(); - assertEquals(1, results.size()); - ProfileResult result = results.get(0); - assertEquals(0, (long) result.getTimeBreakdown().get("build_scorer_count")); - - long rewriteTime = profiler.getRewriteTime(); - assertThat(rewriteTime, greaterThan(0L)); - } - - public void testApproximations() throws IOException { - QueryProfiler profiler = new QueryProfiler(executor != null); - searcher.setProfiler(profiler); - Query query = new RandomApproximationQuery(new TermQuery(new Term("foo", "bar")), random()); - searcher.count(query); - List results = profiler.getTree(); - assertEquals(1, results.size()); - Map breakdown = results.get(0).getTimeBreakdown(); - assertThat(breakdown.get(QueryTimingType.CREATE_WEIGHT.toString()), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.BUILD_SCORER.toString()), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.NEXT_DOC.toString()), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.ADVANCE.toString()), equalTo(0L)); - assertThat(breakdown.get(QueryTimingType.SCORE.toString()), equalTo(0L)); - assertThat(breakdown.get(QueryTimingType.MATCH.toString()), greaterThan(0L)); - - assertThat(breakdown.get(QueryTimingType.CREATE_WEIGHT.toString() + "_count"), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.BUILD_SCORER.toString() + "_count"), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.NEXT_DOC.toString() + "_count"), greaterThan(0L)); - assertThat(breakdown.get(QueryTimingType.ADVANCE.toString() + "_count"), equalTo(0L)); - assertThat(breakdown.get(QueryTimingType.SCORE.toString() + "_count"), equalTo(0L)); - assertThat(breakdown.get(QueryTimingType.MATCH.toString() + "_count"), greaterThan(0L)); - - long rewriteTime = profiler.getRewriteTime(); - assertThat(rewriteTime, greaterThan(0L)); - } - - public void testCollector() throws IOException { - TotalHitCountCollector collector = new TotalHitCountCollector(); - ProfileCollector profileCollector = new ProfileCollector(collector); - assertEquals(0, profileCollector.getTime()); - final LeafCollector leafCollector = profileCollector.getLeafCollector(reader.leaves().get(0)); - assertThat(profileCollector.getTime(), greaterThan(0L)); - long time = profileCollector.getTime(); - leafCollector.setScorer(null); - assertThat(profileCollector.getTime(), greaterThan(time)); - time = profileCollector.getTime(); - leafCollector.collect(0); - assertThat(profileCollector.getTime(), greaterThan(time)); - } - - private static class DummyQuery extends Query { - - @Override - public String toString(String field) { - return getClass().getSimpleName(); - } - - @Override - public boolean equals(Object obj) { - return this == obj; - } - - @Override - public int hashCode() { - return 0; - } - - @Override - public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { - return new Weight(this) { - @Override - public Explanation explain(LeafReaderContext context, int doc) throws IOException { - throw new UnsupportedOperationException(); - } - - @Override - public Scorer scorer(LeafReaderContext context) throws IOException { - throw new UnsupportedOperationException(); - } - - @Override - public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException { - return new ScorerSupplier() { - - @Override - public Scorer get(long loadCost) throws IOException { - throw new UnsupportedOperationException(); - } - - @Override - public long cost() { - return 42; - } - }; - } - - @Override - public boolean isCacheable(LeafReaderContext ctx) { - return true; - } - }; - } - - @Override - public void visit(QueryVisitor visitor) { - visitor.visitLeaf(this); - } - } - - public void testScorerSupplier() throws IOException { - Directory dir = newDirectory(); - IndexWriter w = new IndexWriter(dir, newIndexWriterConfig()); - w.addDocument(new Document()); - DirectoryReader reader = DirectoryReader.open(w); - w.close(); - IndexSearcher s = newSearcher(reader); - s.setQueryCache(null); - Weight weight = s.createWeight(s.rewrite(new DummyQuery()), randomFrom(ScoreMode.values()), 1f); - // exception when getting the scorer - expectThrows(UnsupportedOperationException.class, () -> weight.scorer(s.getIndexReader().leaves().get(0))); - // no exception, means scorerSupplier is delegated - weight.scorerSupplier(s.getIndexReader().leaves().get(0)); - reader.close(); - dir.close(); - } - - private static final QueryCachingPolicy ALWAYS_CACHE_POLICY = new QueryCachingPolicy() { - - @Override - public void onUse(Query query) {} - - @Override - public boolean shouldCache(Query query) throws IOException { - return true; - } - - }; -} diff --git a/sandbox/plugins/concurrent-search/src/test/java/org/opensearch/search/query/QueryPhaseTests.java b/sandbox/plugins/concurrent-search/src/test/java/org/opensearch/search/query/QueryPhaseTests.java deleted file mode 100644 index 2b4ad8f0a26bd..0000000000000 --- a/sandbox/plugins/concurrent-search/src/test/java/org/opensearch/search/query/QueryPhaseTests.java +++ /dev/null @@ -1,1261 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -/* - * Licensed to Elasticsearch under one or more contributor - * license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright - * ownership. Elasticsearch licenses this file to you under - * the Apache License, Version 2.0 (the "License"); you may - * not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -/* - * Modifications Copyright OpenSearch Contributors. See - * GitHub history for details. - */ - -package org.opensearch.search.query; - -import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; -import org.apache.lucene.analysis.standard.StandardAnalyzer; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field.Store; -import org.apache.lucene.document.LatLonDocValuesField; -import org.apache.lucene.document.LatLonPoint; -import org.apache.lucene.document.LongPoint; -import org.apache.lucene.document.NumericDocValuesField; -import org.apache.lucene.document.SortedDocValuesField; -import org.apache.lucene.document.SortedSetDocValuesField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.document.TextField; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.IndexReader; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.index.Term; -import org.apache.lucene.queries.spans.SpanNearQuery; -import org.apache.lucene.queries.spans.SpanTermQuery; -import org.apache.lucene.search.BooleanClause.Occur; -import org.apache.lucene.search.BooleanQuery; -import org.apache.lucene.search.Collector; -import org.apache.lucene.search.ConstantScoreQuery; -import org.apache.lucene.search.DocValuesFieldExistsQuery; -import org.apache.lucene.search.FieldComparator; -import org.apache.lucene.search.FieldDoc; -import org.apache.lucene.search.FilterCollector; -import org.apache.lucene.search.FilterLeafCollector; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.LeafCollector; -import org.apache.lucene.search.MatchAllDocsQuery; -import org.apache.lucene.search.MatchNoDocsQuery; -import org.apache.lucene.search.MultiTermQuery; -import org.apache.lucene.search.PrefixQuery; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.ScoreDoc; -import org.apache.lucene.search.Sort; -import org.apache.lucene.search.SortField; -import org.apache.lucene.search.TermQuery; -import org.apache.lucene.search.TopDocs; -import org.apache.lucene.search.TotalHitCountCollector; -import org.apache.lucene.search.TotalHits; -import org.apache.lucene.search.Weight; -import org.apache.lucene.search.grouping.CollapseTopFieldDocs; -import org.apache.lucene.search.join.BitSetProducer; -import org.apache.lucene.search.join.ScoreMode; -import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.index.RandomIndexWriter; -import org.apache.lucene.util.BytesRef; -import org.apache.lucene.util.FixedBitSet; -import org.opensearch.action.search.SearchShardTask; -import org.opensearch.common.settings.Settings; -import org.opensearch.index.mapper.DateFieldMapper; -import org.opensearch.index.mapper.MappedFieldType; -import org.opensearch.index.mapper.MapperService; -import org.opensearch.index.mapper.NumberFieldMapper; -import org.opensearch.index.mapper.NumberFieldMapper.NumberFieldType; -import org.opensearch.index.mapper.NumberFieldMapper.NumberType; -import org.opensearch.index.query.ParsedQuery; -import org.opensearch.index.query.QueryShardContext; -import org.opensearch.index.search.OpenSearchToParentBlockJoinQuery; -import org.opensearch.index.shard.IndexShard; -import org.opensearch.index.shard.IndexShardTestCase; -import org.opensearch.lucene.queries.MinDocQuery; -import org.opensearch.search.DocValueFormat; -import org.opensearch.search.collapse.CollapseBuilder; -import org.opensearch.search.internal.ContextIndexSearcher; -import org.opensearch.search.internal.ScrollContext; -import org.opensearch.search.internal.SearchContext; -import org.opensearch.search.sort.SortAndFormats; -import org.opensearch.tasks.TaskCancelledException; -import org.opensearch.test.TestSearchContext; -import org.opensearch.threadpool.ThreadPool; - -import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collection; -import java.util.Collections; -import java.util.List; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; - -import static org.hamcrest.Matchers.anyOf; -import static org.hamcrest.Matchers.equalTo; -import static org.hamcrest.Matchers.greaterThan; -import static org.hamcrest.Matchers.greaterThanOrEqualTo; -import static org.hamcrest.Matchers.instanceOf; -import static org.hamcrest.Matchers.lessThanOrEqualTo; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.when; -import static org.opensearch.search.query.TopDocsCollectorContext.hasInfMaxScore; - -public class QueryPhaseTests extends IndexShardTestCase { - - private IndexShard indexShard; - private final ExecutorService executor; - private final QueryPhaseSearcher queryPhaseSearcher; - - @ParametersFactory - public static Collection concurrency() { - return Arrays.asList( - new Object[] { 0, QueryPhase.DEFAULT_QUERY_PHASE_SEARCHER }, - new Object[] { 5, new ConcurrentQueryPhaseSearcher() } - ); - } - - public QueryPhaseTests(int concurrency, QueryPhaseSearcher queryPhaseSearcher) { - this.executor = (concurrency > 0) ? Executors.newFixedThreadPool(concurrency) : null; - this.queryPhaseSearcher = queryPhaseSearcher; - } - - @Override - public Settings threadPoolSettings() { - return Settings.builder().put(super.threadPoolSettings()).put("thread_pool.search.min_queue_size", 10).build(); - } - - @Override - public void setUp() throws Exception { - super.setUp(); - indexShard = newShard(true); - } - - @Override - public void tearDown() throws Exception { - super.tearDown(); - closeShards(indexShard); - - if (executor != null) { - ThreadPool.terminate(executor, 10, TimeUnit.SECONDS); - } - } - - private void countTestCase(Query query, IndexReader reader, boolean shouldCollectSearch, boolean shouldCollectCount) throws Exception { - ContextIndexSearcher searcher = shouldCollectSearch - ? newContextSearcher(reader, executor) - : newEarlyTerminationContextSearcher(reader, 0, executor); - TestSearchContext context = new TestSearchContext(null, indexShard, searcher); - context.parsedQuery(new ParsedQuery(query)); - context.setSize(0); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - final boolean rescore = QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertFalse(rescore); - - ContextIndexSearcher countSearcher = shouldCollectCount - ? newContextSearcher(reader, executor) - : newEarlyTerminationContextSearcher(reader, 0, executor); - assertEquals(countSearcher.count(query), context.queryResult().topDocs().topDocs.totalHits.value); - } - - private void countTestCase(boolean withDeletions) throws Exception { - Directory dir = newDirectory(); - IndexWriterConfig iwc = newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; ++i) { - Document doc = new Document(); - if (randomBoolean()) { - doc.add(new StringField("foo", "bar", Store.NO)); - doc.add(new SortedSetDocValuesField("foo", new BytesRef("bar"))); - doc.add(new SortedSetDocValuesField("docValuesOnlyField", new BytesRef("bar"))); - doc.add(new LatLonDocValuesField("latLonDVField", 1.0, 1.0)); - doc.add(new LatLonPoint("latLonDVField", 1.0, 1.0)); - } - if (randomBoolean()) { - doc.add(new StringField("foo", "baz", Store.NO)); - doc.add(new SortedSetDocValuesField("foo", new BytesRef("baz"))); - } - if (withDeletions && (rarely() || i == 0)) { - doc.add(new StringField("delete", "yes", Store.NO)); - } - w.addDocument(doc); - } - if (withDeletions) { - w.deleteDocuments(new Term("delete", "yes")); - } - final IndexReader reader = w.getReader(); - Query matchAll = new MatchAllDocsQuery(); - Query matchAllCsq = new ConstantScoreQuery(matchAll); - Query tq = new TermQuery(new Term("foo", "bar")); - Query tCsq = new ConstantScoreQuery(tq); - Query dvfeq = new DocValuesFieldExistsQuery("foo"); - Query dvfeq_points = new DocValuesFieldExistsQuery("latLonDVField"); - Query dvfeqCsq = new ConstantScoreQuery(dvfeq); - // field with doc-values but not indexed will need to collect - Query dvOnlyfeq = new DocValuesFieldExistsQuery("docValuesOnlyField"); - BooleanQuery bq = new BooleanQuery.Builder().add(matchAll, Occur.SHOULD).add(tq, Occur.MUST).build(); - - countTestCase(matchAll, reader, false, false); - countTestCase(matchAllCsq, reader, false, false); - countTestCase(tq, reader, withDeletions, withDeletions); - countTestCase(tCsq, reader, withDeletions, withDeletions); - countTestCase(dvfeq, reader, withDeletions, true); - countTestCase(dvfeq_points, reader, withDeletions, true); - countTestCase(dvfeqCsq, reader, withDeletions, true); - countTestCase(dvOnlyfeq, reader, true, true); - countTestCase(bq, reader, true, true); - reader.close(); - w.close(); - dir.close(); - } - - public void testCountWithoutDeletions() throws Exception { - countTestCase(false); - } - - public void testCountWithDeletions() throws Exception { - countTestCase(true); - } - - public void testPostFilterDisablesCountOptimization() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - Document doc = new Document(); - w.addDocument(doc); - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - - TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0, executor)); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value); - - context.setSearcher(newContextSearcher(reader, executor)); - context.parsedPostFilter(new ParsedQuery(new MatchNoDocsQuery())); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertEquals(0, context.queryResult().topDocs().topDocs.totalHits.value); - reader.close(); - dir.close(); - } - - public void testTerminateAfterWithFilter() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - Document doc = new Document(); - for (int i = 0; i < 10; i++) { - doc.add(new StringField("foo", Integer.toString(i), Store.NO)); - } - w.addDocument(doc); - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - context.terminateAfter(1); - context.setSize(10); - for (int i = 0; i < 10; i++) { - context.parsedPostFilter(new ParsedQuery(new TermQuery(new Term("foo", Integer.toString(i))))); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - } - reader.close(); - dir.close(); - } - - public void testMinScoreDisablesCountOptimization() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - Document doc = new Document(); - w.addDocument(doc); - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0, executor)); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - context.setSize(0); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value); - - context.minimumScore(100); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertEquals(0, context.queryResult().topDocs().topDocs.totalHits.value); - assertEquals(TotalHits.Relation.EQUAL_TO, context.queryResult().topDocs().topDocs.totalHits.relation); - reader.close(); - dir.close(); - } - - public void testQueryCapturesThreadPoolStats() throws Exception { - Directory dir = newDirectory(); - IndexWriterConfig iwc = newIndexWriterConfig(); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; ++i) { - w.addDocument(new Document()); - } - w.close(); - IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - QuerySearchResult results = context.queryResult(); - assertThat(results.serviceTimeEWMA(), greaterThanOrEqualTo(0L)); - assertThat(results.nodeQueueSize(), greaterThanOrEqualTo(0)); - reader.close(); - dir.close(); - } - - public void testInOrderScrollOptimization() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; ++i) { - w.addDocument(new Document()); - } - w.close(); - IndexReader reader = DirectoryReader.open(dir); - ScrollContext scrollContext = new ScrollContext(); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor), scrollContext); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - context.sort(new SortAndFormats(sort, new DocValueFormat[] { DocValueFormat.RAW })); - scrollContext.lastEmittedDoc = null; - scrollContext.maxScore = Float.NaN; - scrollContext.totalHits = null; - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - int size = randomIntBetween(2, 5); - context.setSize(size); - - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.terminateAfter(), equalTo(0)); - assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); - - context.setSearcher(newEarlyTerminationContextSearcher(reader, size, executor)); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0].doc, greaterThanOrEqualTo(size)); - reader.close(); - dir.close(); - } - - public void testTerminateAfterEarlyTermination() throws Exception { - Directory dir = newDirectory(); - IndexWriterConfig iwc = newIndexWriterConfig(); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; ++i) { - Document doc = new Document(); - if (randomBoolean()) { - doc.add(new StringField("foo", "bar", Store.NO)); - } - if (randomBoolean()) { - doc.add(new StringField("foo", "baz", Store.NO)); - } - doc.add(new NumericDocValuesField("rank", numDocs - i)); - w.addDocument(doc); - } - w.close(); - final IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - - context.terminateAfter(numDocs); - { - context.setSize(10); - final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(executor); - context.queryCollectorManagers().put(TotalHitCountCollector.class, manager); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertFalse(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(10)); - assertThat(manager.getTotalHits(), equalTo(numDocs)); - } - - context.terminateAfter(1); - { - context.setSize(1); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - - context.setSize(0); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0)); - } - - { - context.setSize(1); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - } - { - context.setSize(1); - BooleanQuery bq = new BooleanQuery.Builder().add(new TermQuery(new Term("foo", "bar")), Occur.SHOULD) - .add(new TermQuery(new Term("foo", "baz")), Occur.SHOULD) - .build(); - context.parsedQuery(new ParsedQuery(bq)); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - - context.setSize(0); - context.parsedQuery(new ParsedQuery(bq)); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0)); - } - { - context.setSize(1); - final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(executor, 1); - context.queryCollectorManagers().put(TotalHitCountCollector.class, manager); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertThat(manager.getTotalHits(), equalTo(1)); - context.queryCollectorManagers().clear(); - } - { - context.setSize(0); - final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(executor, 1); - context.queryCollectorManagers().put(TotalHitCountCollector.class, manager); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0)); - assertThat(manager.getTotalHits(), equalTo(1)); - } - - // tests with trackTotalHits and terminateAfter - context.terminateAfter(10); - context.setSize(0); - for (int trackTotalHits : new int[] { -1, 3, 76, 100 }) { - context.trackTotalHitsUpTo(trackTotalHits); - final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(executor); - context.queryCollectorManagers().put(TotalHitCountCollector.class, manager); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - if (trackTotalHits == -1) { - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(0L)); - } else { - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) Math.min(trackTotalHits, 10))); - } - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0)); - // The concurrent search terminates the collection when the number of hits is reached by each - // concurrent collector. In this case, in general, the number of results are multiplied by the number of - // slices (as the unit of concurrency). To address that, we have to use the shared global state, - // much as HitsThresholdChecker does. - if (executor == null) { - assertThat(manager.getTotalHits(), equalTo(10)); - } - } - - context.terminateAfter(7); - context.setSize(10); - for (int trackTotalHits : new int[] { -1, 3, 75, 100 }) { - context.trackTotalHitsUpTo(trackTotalHits); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - if (trackTotalHits == -1) { - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(0L)); - } else { - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(7L)); - } - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(7)); - } - reader.close(); - dir.close(); - } - - public void testIndexSortingEarlyTermination() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; ++i) { - Document doc = new Document(); - if (randomBoolean()) { - doc.add(new StringField("foo", "bar", Store.NO)); - } - if (randomBoolean()) { - doc.add(new StringField("foo", "baz", Store.NO)); - } - doc.add(new NumericDocValuesField("rank", numDocs - i)); - w.addDocument(doc); - } - w.close(); - - final IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - context.setSize(1); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.sort(new SortAndFormats(sort, new DocValueFormat[] { DocValueFormat.RAW })); - - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0], instanceOf(FieldDoc.class)); - FieldDoc fieldDoc = (FieldDoc) context.queryResult().topDocs().topDocs.scoreDocs[0]; - assertThat(fieldDoc.fields[0], equalTo(1)); - - { - context.parsedPostFilter(new ParsedQuery(new MinDocQuery(1))); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(numDocs - 1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0], instanceOf(FieldDoc.class)); - assertThat(fieldDoc.fields[0], anyOf(equalTo(1), equalTo(2))); - context.parsedPostFilter(null); - - final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(executor, sort); - context.queryCollectorManagers().put(TotalHitCountCollector.class, manager); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0], instanceOf(FieldDoc.class)); - assertThat(fieldDoc.fields[0], anyOf(equalTo(1), equalTo(2))); - // When searching concurrently, each executors short-circuits when "size" is reached, - // including total hits collector - assertThat(manager.getTotalHits(), lessThanOrEqualTo(numDocs)); - - context.queryCollectorManagers().clear(); - } - - { - context.setSearcher(newEarlyTerminationContextSearcher(reader, 1, executor)); - context.trackTotalHitsUpTo(SearchContext.TRACK_TOTAL_HITS_DISABLED); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0], instanceOf(FieldDoc.class)); - assertThat(fieldDoc.fields[0], anyOf(equalTo(1), equalTo(2))); - - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0], instanceOf(FieldDoc.class)); - assertThat(fieldDoc.fields[0], anyOf(equalTo(1), equalTo(2))); - } - reader.close(); - dir.close(); - } - - public void testIndexSortScrollOptimization() throws Exception { - Directory dir = newDirectory(); - final Sort indexSort = new Sort(new SortField("rank", SortField.Type.INT), new SortField("tiebreaker", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(indexSort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; ++i) { - Document doc = new Document(); - doc.add(new NumericDocValuesField("rank", random().nextInt())); - doc.add(new NumericDocValuesField("tiebreaker", i)); - w.addDocument(doc); - } - if (randomBoolean()) { - w.forceMerge(randomIntBetween(1, 10)); - } - w.close(); - - final IndexReader reader = DirectoryReader.open(dir); - List searchSortAndFormats = new ArrayList<>(); - searchSortAndFormats.add(new SortAndFormats(indexSort, new DocValueFormat[] { DocValueFormat.RAW, DocValueFormat.RAW })); - // search sort is a prefix of the index sort - searchSortAndFormats.add(new SortAndFormats(new Sort(indexSort.getSort()[0]), new DocValueFormat[] { DocValueFormat.RAW })); - for (SortAndFormats searchSortAndFormat : searchSortAndFormats) { - ScrollContext scrollContext = new ScrollContext(); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor), scrollContext); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - scrollContext.lastEmittedDoc = null; - scrollContext.maxScore = Float.NaN; - scrollContext.totalHits = null; - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.setSize(10); - context.sort(searchSortAndFormat); - - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.terminateAfter(), equalTo(0)); - assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); - int sizeMinus1 = context.queryResult().topDocs().topDocs.scoreDocs.length - 1; - FieldDoc lastDoc = (FieldDoc) context.queryResult().topDocs().topDocs.scoreDocs[sizeMinus1]; - - context.setSearcher(newEarlyTerminationContextSearcher(reader, 10, executor)); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.terminateAfter(), equalTo(0)); - assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); - FieldDoc firstDoc = (FieldDoc) context.queryResult().topDocs().topDocs.scoreDocs[0]; - for (int i = 0; i < searchSortAndFormat.sort.getSort().length; i++) { - @SuppressWarnings("unchecked") - FieldComparator comparator = (FieldComparator) searchSortAndFormat.sort.getSort()[i].getComparator( - i, - false - ); - int cmp = comparator.compareValues(firstDoc.fields[i], lastDoc.fields[i]); - if (cmp == 0) { - continue; - } - assertThat(cmp, equalTo(1)); - break; - } - } - reader.close(); - dir.close(); - } - - public void testDisableTopScoreCollection() throws Exception { - Directory dir = newDirectory(); - IndexWriterConfig iwc = newIndexWriterConfig(new StandardAnalyzer()); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - Document doc = new Document(); - final int numDocs = 2 * scaledRandomIntBetween(50, 450); - for (int i = 0; i < numDocs; i++) { - doc.clear(); - if (i % 2 == 0) { - doc.add(new TextField("title", "foo bar", Store.NO)); - } else { - doc.add(new TextField("title", "foo", Store.NO)); - } - w.addDocument(doc); - } - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - Query q = new SpanNearQuery.Builder("title", true).addClause(new SpanTermQuery(new Term("title", "foo"))) - .addClause(new SpanTermQuery(new Term("title", "bar"))) - .build(); - - context.parsedQuery(new ParsedQuery(q)); - context.setSize(3); - context.trackTotalHitsUpTo(3); - TopDocsCollectorContext topDocsContext = TopDocsCollectorContext.createTopDocsCollectorContext(context, false); - assertEquals(topDocsContext.create(null).scoreMode(), org.apache.lucene.search.ScoreMode.COMPLETE); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertEquals(numDocs / 2, context.queryResult().topDocs().topDocs.totalHits.value); - assertEquals(context.queryResult().topDocs().topDocs.totalHits.relation, TotalHits.Relation.EQUAL_TO); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(3)); - - 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); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertEquals(numDocs / 2, context.queryResult().topDocs().topDocs.totalHits.value); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(3)); - assertEquals(context.queryResult().topDocs().topDocs.totalHits.relation, TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO); - - reader.close(); - dir.close(); - } - - public void testEnhanceSortOnNumeric() throws Exception { - final String fieldNameLong = "long-field"; - final String fieldNameDate = "date-field"; - MappedFieldType fieldTypeLong = new NumberFieldMapper.NumberFieldType(fieldNameLong, NumberFieldMapper.NumberType.LONG); - MappedFieldType fieldTypeDate = new DateFieldMapper.DateFieldType(fieldNameDate); - MapperService mapperService = mock(MapperService.class); - when(mapperService.fieldType(fieldNameLong)).thenReturn(fieldTypeLong); - when(mapperService.fieldType(fieldNameDate)).thenReturn(fieldTypeDate); - // enough docs to have a tree with several leaf nodes - final int numDocs = 3500 * 5; - Directory dir = newDirectory(); - IndexWriter writer = new IndexWriter(dir, new IndexWriterConfig(null)); - long firstValue = randomLongBetween(-10000000L, 10000000L); - long longValue = firstValue; - long dateValue = randomLongBetween(0, 3000000000000L); - for (int i = 1; i <= numDocs; ++i) { - Document doc = new Document(); - - doc.add(new LongPoint(fieldNameLong, longValue)); - doc.add(new NumericDocValuesField(fieldNameLong, longValue)); - - doc.add(new LongPoint(fieldNameDate, dateValue)); - doc.add(new NumericDocValuesField(fieldNameDate, dateValue)); - writer.addDocument(doc); - longValue++; - dateValue++; - if (i % 3500 == 0) writer.commit(); - } - writer.close(); - final IndexReader reader = DirectoryReader.open(dir); - final SortField sortFieldLong = new SortField(fieldNameLong, SortField.Type.LONG); - sortFieldLong.setMissingValue(Long.MAX_VALUE); - final SortField sortFieldDate = new SortField(fieldNameDate, SortField.Type.LONG); - sortFieldDate.setMissingValue(Long.MAX_VALUE); - DocValueFormat dateFormat = fieldTypeDate.docValueFormat(null, null); - final Sort longSort = new Sort(sortFieldLong); - final Sort longDateSort = new Sort(sortFieldLong, sortFieldDate); - final Sort dateSort = new Sort(sortFieldDate); - final Sort dateLongSort = new Sort(sortFieldDate, sortFieldLong); - SortAndFormats longSortAndFormats = new SortAndFormats(longSort, new DocValueFormat[] { DocValueFormat.RAW }); - SortAndFormats longDateSortAndFormats = new SortAndFormats(longDateSort, new DocValueFormat[] { DocValueFormat.RAW, dateFormat }); - SortAndFormats dateSortAndFormats = new SortAndFormats(dateSort, new DocValueFormat[] { dateFormat }); - SortAndFormats dateLongSortAndFormats = new SortAndFormats(dateLongSort, new DocValueFormat[] { dateFormat, DocValueFormat.RAW }); - ParsedQuery query = new ParsedQuery(new MatchAllDocsQuery()); - SearchShardTask task = new SearchShardTask(123L, "", "", "", null, Collections.emptyMap()); - - // 1. Test a sort on long field - { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); - when(searchContext.mapperService()).thenReturn(mapperService); - searchContext.sort(longSortAndFormats); - searchContext.parsedQuery(query); - searchContext.setTask(task); - searchContext.setSize(10); - QueryPhase.executeInternal(searchContext.withCleanQueryResult(), queryPhaseSearcher); - assertSortResults(searchContext.queryResult().topDocs().topDocs, (long) numDocs, false); - } - - // 2. Test a sort on long field + date field - { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); - when(searchContext.mapperService()).thenReturn(mapperService); - searchContext.sort(longDateSortAndFormats); - searchContext.parsedQuery(query); - searchContext.setTask(task); - searchContext.setSize(10); - QueryPhase.executeInternal(searchContext.withCleanQueryResult(), queryPhaseSearcher); - assertSortResults(searchContext.queryResult().topDocs().topDocs, (long) numDocs, true); - } - - // 3. Test a sort on date field - { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); - when(searchContext.mapperService()).thenReturn(mapperService); - searchContext.sort(dateSortAndFormats); - searchContext.parsedQuery(query); - searchContext.setTask(task); - searchContext.setSize(10); - QueryPhase.executeInternal(searchContext.withCleanQueryResult(), queryPhaseSearcher); - assertSortResults(searchContext.queryResult().topDocs().topDocs, (long) numDocs, false); - } - - // 4. Test a sort on date field + long field - { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); - when(searchContext.mapperService()).thenReturn(mapperService); - searchContext.sort(dateLongSortAndFormats); - searchContext.parsedQuery(query); - searchContext.setTask(task); - searchContext.setSize(10); - QueryPhase.executeInternal(searchContext.withCleanQueryResult(), queryPhaseSearcher); - assertSortResults(searchContext.queryResult().topDocs().topDocs, (long) numDocs, true); - } - - // 5. Test that sort optimization is run when from > 0 and size = 0 - { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); - when(searchContext.mapperService()).thenReturn(mapperService); - searchContext.sort(longSortAndFormats); - searchContext.parsedQuery(query); - searchContext.setTask(task); - searchContext.from(5); - searchContext.setSize(0); - QueryPhase.executeInternal(searchContext.withCleanQueryResult(), queryPhaseSearcher); - assertSortResults(searchContext.queryResult().topDocs().topDocs, (long) numDocs, false); - } - - // 6. Test that sort optimization works with from = 0 and size= 0 - { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); - when(searchContext.mapperService()).thenReturn(mapperService); - searchContext.sort(longSortAndFormats); - searchContext.parsedQuery(query); - searchContext.setTask(task); - searchContext.setSize(0); - QueryPhase.executeInternal(searchContext.withCleanQueryResult(), queryPhaseSearcher); - } - - // 7. Test that sort optimization works with search after - { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); - when(searchContext.mapperService()).thenReturn(mapperService); - int afterDocument = (int) randomLongBetween(0, 50); - long afterValue = firstValue + afterDocument; - FieldDoc after = new FieldDoc(afterDocument, Float.NaN, new Long[] { afterValue }); - searchContext.searchAfter(after); - searchContext.sort(longSortAndFormats); - searchContext.parsedQuery(query); - searchContext.setTask(task); - searchContext.setSize(10); - QueryPhase.executeInternal(searchContext.withCleanQueryResult(), queryPhaseSearcher); - final TopDocs topDocs = searchContext.queryResult().topDocs().topDocs; - long topValue = (long) ((FieldDoc) topDocs.scoreDocs[0]).fields[0]; - assertThat(topValue, greaterThan(afterValue)); - assertSortResults(topDocs, (long) numDocs, false); - - final TotalHits totalHits = topDocs.totalHits; - assertEquals(TotalHits.Relation.EQUAL_TO, totalHits.relation); - assertEquals(numDocs, totalHits.value); - } - - reader.close(); - dir.close(); - } - - public void testMaxScoreQueryVisitor() { - BitSetProducer producer = context -> new FixedBitSet(1); - Query query = new OpenSearchToParentBlockJoinQuery(new MatchAllDocsQuery(), producer, ScoreMode.Avg, "nested"); - assertTrue(hasInfMaxScore(query)); - - query = new OpenSearchToParentBlockJoinQuery(new MatchAllDocsQuery(), producer, ScoreMode.None, "nested"); - assertFalse(hasInfMaxScore(query)); - - for (Occur occur : Occur.values()) { - query = new BooleanQuery.Builder().add( - new OpenSearchToParentBlockJoinQuery(new MatchAllDocsQuery(), producer, ScoreMode.Avg, "nested"), - occur - ).build(); - if (occur == Occur.MUST) { - assertTrue(hasInfMaxScore(query)); - } else { - assertFalse(hasInfMaxScore(query)); - } - - query = new BooleanQuery.Builder().add( - new BooleanQuery.Builder().add( - new OpenSearchToParentBlockJoinQuery(new MatchAllDocsQuery(), producer, ScoreMode.Avg, "nested"), - occur - ).build(), - occur - ).build(); - if (occur == Occur.MUST) { - assertTrue(hasInfMaxScore(query)); - } else { - assertFalse(hasInfMaxScore(query)); - } - - query = new BooleanQuery.Builder().add( - new BooleanQuery.Builder().add( - new OpenSearchToParentBlockJoinQuery(new MatchAllDocsQuery(), producer, ScoreMode.Avg, "nested"), - occur - ).build(), - Occur.FILTER - ).build(); - assertFalse(hasInfMaxScore(query)); - - query = new BooleanQuery.Builder().add( - new BooleanQuery.Builder().add(new SpanTermQuery(new Term("field", "foo")), occur) - .add(new OpenSearchToParentBlockJoinQuery(new MatchAllDocsQuery(), producer, ScoreMode.Avg, "nested"), occur) - .build(), - occur - ).build(); - if (occur == Occur.MUST) { - assertTrue(hasInfMaxScore(query)); - } else { - assertFalse(hasInfMaxScore(query)); - } - } - } - - // assert score docs are in order and their number is as expected - private void assertSortResults(TopDocs topDocs, long expectedNumDocs, boolean isDoubleSort) { - if (topDocs.totalHits.relation == TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO) { - assertThat(topDocs.totalHits.value, lessThanOrEqualTo(expectedNumDocs)); - } else { - assertEquals(topDocs.totalHits.value, expectedNumDocs); - } - long cur1, cur2; - long prev1 = Long.MIN_VALUE; - long prev2 = Long.MIN_VALUE; - for (ScoreDoc scoreDoc : topDocs.scoreDocs) { - cur1 = (long) ((FieldDoc) scoreDoc).fields[0]; - assertThat(cur1, greaterThanOrEqualTo(prev1)); // test that docs are properly sorted on the first sort - if (isDoubleSort) { - cur2 = (long) ((FieldDoc) scoreDoc).fields[1]; - if (cur1 == prev1) { - assertThat(cur2, greaterThanOrEqualTo(prev2)); // test that docs are properly sorted on the secondary sort - } - prev2 = cur2; - } - prev1 = cur1; - } - } - - public void testMinScore() throws Exception { - Directory dir = newDirectory(); - IndexWriterConfig iwc = newIndexWriterConfig(); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - for (int i = 0; i < 10; i++) { - Document doc = new Document(); - doc.add(new StringField("foo", "bar", Store.NO)); - doc.add(new StringField("filter", "f1", Store.NO)); - w.addDocument(doc); - } - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.parsedQuery( - new ParsedQuery( - new BooleanQuery.Builder().add(new TermQuery(new Term("foo", "bar")), Occur.MUST) - .add(new TermQuery(new Term("filter", "f1")), Occur.SHOULD) - .build() - ) - ); - context.minimumScore(0.01f); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.setSize(1); - context.trackTotalHitsUpTo(5); - - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertEquals(10, context.queryResult().topDocs().topDocs.totalHits.value); - - reader.close(); - dir.close(); - } - - public void testMaxScore() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("filter", SortField.Type.STRING)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; i++) { - Document doc = new Document(); - doc.add(new StringField("foo", "bar", Store.NO)); - doc.add(new StringField("filter", "f1" + ((i > 0) ? " " + Integer.toString(i) : ""), Store.NO)); - doc.add(new SortedDocValuesField("filter", newBytesRef("f1" + ((i > 0) ? " " + Integer.toString(i) : "")))); - w.addDocument(doc); - } - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.trackScores(true); - context.parsedQuery( - new ParsedQuery( - new BooleanQuery.Builder().add(new TermQuery(new Term("foo", "bar")), Occur.MUST) - .add(new TermQuery(new Term("filter", "f1")), Occur.SHOULD) - .build() - ) - ); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.setSize(1); - context.trackTotalHitsUpTo(5); - - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertFalse(Float.isNaN(context.queryResult().getMaxScore())); - assertEquals(1, context.queryResult().topDocs().topDocs.scoreDocs.length); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, greaterThanOrEqualTo(6L)); - - context.sort(new SortAndFormats(sort, new DocValueFormat[] { DocValueFormat.RAW })); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertFalse(Float.isNaN(context.queryResult().getMaxScore())); - assertEquals(1, context.queryResult().topDocs().topDocs.scoreDocs.length); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, greaterThanOrEqualTo(6L)); - - context.trackScores(false); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(Float.isNaN(context.queryResult().getMaxScore())); - assertEquals(1, context.queryResult().topDocs().topDocs.scoreDocs.length); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, greaterThanOrEqualTo(6L)); - - reader.close(); - dir.close(); - } - - public void testCollapseQuerySearchResults() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("user", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - - // Always end up with uneven buckets so collapsing is predictable - final int numDocs = 2 * scaledRandomIntBetween(600, 900) - 1; - for (int i = 0; i < numDocs; i++) { - Document doc = new Document(); - doc.add(new StringField("foo", "bar", Store.NO)); - doc.add(new NumericDocValuesField("user", i & 1)); - w.addDocument(doc); - } - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - QueryShardContext queryShardContext = mock(QueryShardContext.class); - when(queryShardContext.fieldMapper("user")).thenReturn( - new NumberFieldType("user", NumberType.INTEGER, true, false, true, false, null, Collections.emptyMap()) - ); - - TestSearchContext context = new TestSearchContext(queryShardContext, indexShard, newContextSearcher(reader, executor)); - context.collapse(new CollapseBuilder("user").build(context.getQueryShardContext())); - context.trackScores(true); - context.parsedQuery(new ParsedQuery(new TermQuery(new Term("foo", "bar")))); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.setSize(2); - context.trackTotalHitsUpTo(5); - - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertFalse(Float.isNaN(context.queryResult().getMaxScore())); - assertEquals(2, context.queryResult().topDocs().topDocs.scoreDocs.length); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs, instanceOf(CollapseTopFieldDocs.class)); - - CollapseTopFieldDocs topDocs = (CollapseTopFieldDocs) context.queryResult().topDocs().topDocs; - assertThat(topDocs.collapseValues.length, equalTo(2)); - assertThat(topDocs.collapseValues[0], equalTo(0L)); // user == 0 - assertThat(topDocs.collapseValues[1], equalTo(1L)); // user == 1 - - context.sort(new SortAndFormats(sort, new DocValueFormat[] { DocValueFormat.RAW })); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertFalse(Float.isNaN(context.queryResult().getMaxScore())); - assertEquals(2, context.queryResult().topDocs().topDocs.scoreDocs.length); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs, instanceOf(CollapseTopFieldDocs.class)); - - topDocs = (CollapseTopFieldDocs) context.queryResult().topDocs().topDocs; - assertThat(topDocs.collapseValues.length, equalTo(2)); - assertThat(topDocs.collapseValues[0], equalTo(0L)); // user == 0 - assertThat(topDocs.collapseValues[1], equalTo(1L)); // user == 1 - - context.trackScores(false); - QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher); - assertTrue(Float.isNaN(context.queryResult().getMaxScore())); - assertEquals(2, context.queryResult().topDocs().topDocs.scoreDocs.length); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs, instanceOf(CollapseTopFieldDocs.class)); - - topDocs = (CollapseTopFieldDocs) context.queryResult().topDocs().topDocs; - assertThat(topDocs.collapseValues.length, equalTo(2)); - assertThat(topDocs.collapseValues[0], equalTo(0L)); // user == 0 - assertThat(topDocs.collapseValues[1], equalTo(1L)); // user == 1 - - reader.close(); - dir.close(); - } - - public void testCancellationDuringPreprocess() throws IOException { - try (Directory dir = newDirectory(); RandomIndexWriter w = new RandomIndexWriter(random(), dir, newIndexWriterConfig())) { - - for (int i = 0; i < 10; i++) { - Document doc = new Document(); - StringBuilder sb = new StringBuilder(); - for (int j = 0; j < i; j++) { - sb.append('a'); - } - doc.add(new StringField("foo", sb.toString(), Store.NO)); - w.addDocument(doc); - } - w.flush(); - w.close(); - - try (IndexReader reader = DirectoryReader.open(dir)) { - TestSearchContext context = new TestSearchContextWithRewriteAndCancellation( - null, - indexShard, - newContextSearcher(reader, executor) - ); - PrefixQuery prefixQuery = new PrefixQuery(new Term("foo", "a"), MultiTermQuery.SCORING_BOOLEAN_REWRITE); - context.parsedQuery(new ParsedQuery(prefixQuery)); - SearchShardTask task = mock(SearchShardTask.class); - when(task.isCancelled()).thenReturn(true); - context.setTask(task); - expectThrows(TaskCancelledException.class, () -> new QueryPhase().preProcess(context)); - } - } - } - - private static class TestSearchContextWithRewriteAndCancellation extends TestSearchContext { - - private TestSearchContextWithRewriteAndCancellation( - QueryShardContext queryShardContext, - IndexShard indexShard, - ContextIndexSearcher searcher - ) { - super(queryShardContext, indexShard, searcher); - } - - @Override - public void preProcess(boolean rewrite) { - try { - searcher().rewrite(query()); - } catch (IOException e) { - fail("IOException shouldn't be thrown"); - } - } - - @Override - public boolean lowLevelCancellation() { - return true; - } - } - - private static ContextIndexSearcher newContextSearcher(IndexReader reader, ExecutorService executor) throws IOException { - return new ContextIndexSearcher( - reader, - IndexSearcher.getDefaultSimilarity(), - IndexSearcher.getDefaultQueryCache(), - IndexSearcher.getDefaultQueryCachingPolicy(), - true, - executor - ); - } - - private static ContextIndexSearcher newEarlyTerminationContextSearcher(IndexReader reader, int size, ExecutorService executor) - throws IOException { - return new ContextIndexSearcher( - reader, - IndexSearcher.getDefaultSimilarity(), - IndexSearcher.getDefaultQueryCache(), - IndexSearcher.getDefaultQueryCachingPolicy(), - true, - executor - ) { - - @Override - public void search(List leaves, Weight weight, Collector collector) throws IOException { - final Collector in = new AssertingEarlyTerminationFilterCollector(collector, size); - super.search(leaves, weight, in); - } - }; - } - - private static class TestTotalHitCountCollectorManager extends TotalHitCountCollectorManager { - private int totalHits; - private final TotalHitCountCollector collector; - private final Integer teminateAfter; - - static TestTotalHitCountCollectorManager create(final ExecutorService executor) { - return create(executor, null, null); - } - - static TestTotalHitCountCollectorManager create(final ExecutorService executor, final Integer teminateAfter) { - return create(executor, null, teminateAfter); - } - - static TestTotalHitCountCollectorManager create(final ExecutorService executor, final Sort sort) { - return create(executor, sort, null); - } - - static TestTotalHitCountCollectorManager create(final ExecutorService executor, final Sort sort, final Integer teminateAfter) { - if (executor == null) { - return new TestTotalHitCountCollectorManager(new TotalHitCountCollector(), sort); - } else { - return new TestTotalHitCountCollectorManager(sort, teminateAfter); - } - } - - private TestTotalHitCountCollectorManager(final TotalHitCountCollector collector, final Sort sort) { - super(sort); - this.collector = collector; - this.teminateAfter = null; - } - - private TestTotalHitCountCollectorManager(final Sort sort, final Integer teminateAfter) { - super(sort); - this.collector = null; - this.teminateAfter = teminateAfter; - } - - @Override - public TotalHitCountCollector newCollector() throws IOException { - return (collector == null) ? super.newCollector() : collector; - } - - @Override - public ReduceableSearchResult reduce(Collection collectors) throws IOException { - final ReduceableSearchResult result = super.reduce(collectors); - totalHits = collectors.stream().mapToInt(TotalHitCountCollector::getTotalHits).sum(); - - if (teminateAfter != null) { - assertThat(totalHits, greaterThanOrEqualTo(teminateAfter)); - totalHits = Math.min(totalHits, teminateAfter); - } - - return result; - } - - public int getTotalHits() { - return (collector == null) ? totalHits : collector.getTotalHits(); - } - } - - private static class AssertingEarlyTerminationFilterCollector extends FilterCollector { - private final int size; - - AssertingEarlyTerminationFilterCollector(Collector in, int size) { - super(in); - this.size = size; - } - - @Override - public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException { - final LeafCollector in = super.getLeafCollector(context); - return new FilterLeafCollector(in) { - int collected; - - @Override - public void collect(int doc) throws IOException { - assert collected <= size : "should not collect more than " + size + " doc per segment, got " + collected; - ++collected; - super.collect(doc); - } - }; - } - } -} diff --git a/sandbox/plugins/concurrent-search/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java b/sandbox/plugins/concurrent-search/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java deleted file mode 100644 index ed10fe5cd84ca..0000000000000 --- a/sandbox/plugins/concurrent-search/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java +++ /dev/null @@ -1,1213 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.search.query; - -import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; - -import org.apache.lucene.analysis.standard.StandardAnalyzer; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field.Store; -import org.apache.lucene.document.NumericDocValuesField; -import org.apache.lucene.document.SortedDocValuesField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.document.TextField; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.IndexReader; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.Term; -import org.apache.lucene.queries.spans.SpanNearQuery; -import org.apache.lucene.queries.spans.SpanTermQuery; -import org.apache.lucene.search.BooleanClause.Occur; -import org.apache.lucene.search.grouping.CollapseTopFieldDocs; -import org.apache.lucene.search.BooleanQuery; -import org.apache.lucene.search.Collector; -import org.apache.lucene.search.FieldComparator; -import org.apache.lucene.search.FieldDoc; -import org.apache.lucene.search.FilterCollector; -import org.apache.lucene.search.FilterLeafCollector; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.LeafCollector; -import org.apache.lucene.search.MatchAllDocsQuery; -import org.apache.lucene.search.MatchNoDocsQuery; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.Sort; -import org.apache.lucene.search.SortField; -import org.apache.lucene.search.TermQuery; -import org.apache.lucene.search.TotalHits; -import org.apache.lucene.search.Weight; -import org.apache.lucene.tests.index.RandomIndexWriter; -import org.apache.lucene.store.Directory; -import org.opensearch.action.search.SearchShardTask; -import org.opensearch.common.settings.Settings; -import org.opensearch.core.xcontent.ToXContent; -import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.common.xcontent.json.JsonXContent; -import org.opensearch.index.mapper.NumberFieldMapper.NumberFieldType; -import org.opensearch.index.mapper.NumberFieldMapper.NumberType; -import org.opensearch.index.query.ParsedQuery; -import org.opensearch.index.query.QueryShardContext; -import org.opensearch.index.shard.IndexShard; -import org.opensearch.index.shard.IndexShardTestCase; -import org.opensearch.lucene.queries.MinDocQuery; -import org.opensearch.search.DocValueFormat; -import org.opensearch.search.collapse.CollapseBuilder; -import org.opensearch.search.internal.ContextIndexSearcher; -import org.opensearch.search.internal.ScrollContext; -import org.opensearch.search.internal.SearchContext; -import org.opensearch.search.profile.ProfileResult; -import org.opensearch.search.profile.ProfileShardResult; -import org.opensearch.search.profile.SearchProfileShardResults; -import org.opensearch.search.profile.query.CollectorResult; -import org.opensearch.search.profile.query.QueryProfileShardResult; -import org.opensearch.search.sort.SortAndFormats; -import org.opensearch.test.TestSearchContext; -import org.opensearch.threadpool.ThreadPool; - -import java.io.ByteArrayOutputStream; -import java.io.IOException; -import java.io.OutputStream; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collection; -import java.util.Collections; -import java.util.List; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import java.util.function.Consumer; - -import static org.hamcrest.CoreMatchers.not; -import static org.hamcrest.CoreMatchers.nullValue; -import static org.hamcrest.Matchers.anyOf; -import static org.hamcrest.Matchers.equalTo; -import static org.hamcrest.Matchers.empty; -import static org.hamcrest.Matchers.greaterThan; -import static org.hamcrest.Matchers.greaterThanOrEqualTo; -import static org.hamcrest.Matchers.instanceOf; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; -import static org.hamcrest.Matchers.hasSize; - -public class QueryProfilePhaseTests extends IndexShardTestCase { - - private IndexShard indexShard; - private final ExecutorService executor; - private final QueryPhaseSearcher queryPhaseSearcher; - - @ParametersFactory - public static Collection concurrency() { - return Arrays.asList( - new Object[] { 0, QueryPhase.DEFAULT_QUERY_PHASE_SEARCHER }, - new Object[] { 5, new ConcurrentQueryPhaseSearcher() } - ); - } - - public QueryProfilePhaseTests(int concurrency, QueryPhaseSearcher queryPhaseSearcher) { - this.executor = (concurrency > 0) ? Executors.newFixedThreadPool(concurrency) : null; - this.queryPhaseSearcher = queryPhaseSearcher; - } - - @Override - public Settings threadPoolSettings() { - return Settings.builder().put(super.threadPoolSettings()).put("thread_pool.search.min_queue_size", 10).build(); - } - - @Override - public void setUp() throws Exception { - super.setUp(); - indexShard = newShard(true); - } - - @Override - public void tearDown() throws Exception { - super.tearDown(); - closeShards(indexShard); - - if (executor != null) { - ThreadPool.terminate(executor, 10, TimeUnit.SECONDS); - } - } - - public void testPostFilterDisablesCountOptimization() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - Document doc = new Document(); - w.addDocument(doc); - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - - TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0, executor)); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value); - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_count")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - context.setSearcher(newContextSearcher(reader, executor)); - context.parsedPostFilter(new ParsedQuery(new MatchNoDocsQuery())); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertEquals(0, context.queryResult().topDocs().topDocs.totalHits.value); - assertProfileData(context, collector -> { - assertThat(collector.getReason(), equalTo("search_post_filter")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_count")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }, (query) -> { - assertThat(query.getQueryName(), equalTo("MatchNoDocsQuery")); - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, (query) -> { - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertThat(query.getQueryName(), equalTo("ConstantScoreQuery")); - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }); - - reader.close(); - dir.close(); - } - - public void testTerminateAfterWithFilter() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - Document doc = new Document(); - for (int i = 0; i < 10; i++) { - doc.add(new StringField("foo", Integer.toString(i), Store.NO)); - } - w.addDocument(doc); - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - context.terminateAfter(1); - context.setSize(10); - for (int i = 0; i < 10; i++) { - context.parsedPostFilter(new ParsedQuery(new TermQuery(new Term("foo", Integer.toString(i))))); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertProfileData(context, collector -> { - assertThat(collector.getReason(), equalTo("search_post_filter")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_terminate_after_count")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren().get(0).getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getProfiledChildren().get(0).getReason(), equalTo("search_top_hits")); - assertThat(collector.getProfiledChildren().get(0).getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }, (query) -> { - assertThat(query.getQueryName(), equalTo("TermQuery")); - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, (query) -> { - assertThat(query.getQueryName(), equalTo("MatchAllDocsQuery")); - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(1L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }); - } - reader.close(); - dir.close(); - } - - public void testMinScoreDisablesCountOptimization() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - Document doc = new Document(); - w.addDocument(doc); - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0, executor)); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - context.setSize(0); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value); - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_count")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - context.minimumScore(100); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertEquals(0, context.queryResult().topDocs().topDocs.totalHits.value); - assertEquals(TotalHits.Relation.EQUAL_TO, context.queryResult().topDocs().topDocs.totalHits.relation); - assertProfileData(context, "MatchAllDocsQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThanOrEqualTo(100L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(1L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_min_score")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_count")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }); - - reader.close(); - dir.close(); - } - - public void testInOrderScrollOptimization() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; ++i) { - w.addDocument(new Document()); - } - w.close(); - IndexReader reader = DirectoryReader.open(dir); - ScrollContext scrollContext = new ScrollContext(); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor), scrollContext); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - context.sort(new SortAndFormats(sort, new DocValueFormat[] { DocValueFormat.RAW })); - scrollContext.lastEmittedDoc = null; - scrollContext.maxScore = Float.NaN; - scrollContext.totalHits = null; - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - int size = randomIntBetween(2, 5); - context.setSize(size); - - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.terminateAfter(), equalTo(0)); - assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - context.setSearcher(newEarlyTerminationContextSearcher(reader, size, executor)); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0].doc, greaterThanOrEqualTo(size)); - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(0)); - }); - - reader.close(); - dir.close(); - } - - public void testTerminateAfterEarlyTermination() throws Exception { - Directory dir = newDirectory(); - IndexWriterConfig iwc = newIndexWriterConfig(); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; ++i) { - Document doc = new Document(); - if (randomBoolean()) { - doc.add(new StringField("foo", "bar", Store.NO)); - } - if (randomBoolean()) { - doc.add(new StringField("foo", "baz", Store.NO)); - } - doc.add(new NumericDocValuesField("rank", numDocs - i)); - w.addDocument(doc); - } - w.close(); - final IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - - context.terminateAfter(1); - { - context.setSize(1); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertProfileData(context, "MatchAllDocsQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_terminate_after_count")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_top_hits")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }); - - context.setSize(0); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0)); - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_terminate_after_count")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_count")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }); - } - - { - context.setSize(1); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertProfileData(context, "MatchAllDocsQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_terminate_after_count")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_top_hits")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }); - } - { - context.setSize(1); - BooleanQuery bq = new BooleanQuery.Builder().add(new TermQuery(new Term("foo", "bar")), Occur.SHOULD) - .add(new TermQuery(new Term("foo", "baz")), Occur.SHOULD) - .build(); - context.parsedQuery(new ParsedQuery(bq)); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertProfileData(context, "BooleanQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - assertThat(query.getProfiledChildren(), hasSize(2)); - assertThat(query.getProfiledChildren().get(0).getQueryName(), equalTo("TermQuery")); - assertThat(query.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - assertThat(query.getProfiledChildren().get(1).getQueryName(), equalTo("TermQuery")); - assertThat(query.getProfiledChildren().get(1).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_terminate_after_count")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_top_hits")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }); - context.setSize(0); - context.parsedQuery(new ParsedQuery(bq)); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0)); - - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - // rewritten as a ConstantScoreQuery wrapping the original BooleanQuery - // see: https://github.com/apache/lucene/pull/672 - assertThat(query.getProfiledChildren(), hasSize(1)); - assertThat(query.getProfiledChildren().get(0).getQueryName(), equalTo("BooleanQuery")); - assertThat(query.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("score_count"), equalTo(0L)); - - List children = query.getProfiledChildren().get(0).getProfiledChildren(); - assertThat(children, hasSize(2)); - assertThat(children.get(0).getQueryName(), equalTo("TermQuery")); - assertThat(children.get(0).getTime(), greaterThan(0L)); - assertThat(children.get(0).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(children.get(0).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - assertThat(children.get(0).getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(children.get(0).getTimeBreakdown().get("score_count"), equalTo(0L)); - - assertThat(children.get(1).getQueryName(), equalTo("TermQuery")); - assertThat(children.get(1).getTime(), greaterThan(0L)); - assertThat(children.get(1).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(children.get(1).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - assertThat(children.get(1).getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(children.get(1).getTimeBreakdown().get("score_count"), equalTo(0L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_terminate_after_count")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_count")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }); - } - - context.terminateAfter(7); - context.setSize(10); - for (int trackTotalHits : new int[] { -1, 3, 75, 100 }) { - context.trackTotalHitsUpTo(trackTotalHits); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertTrue(context.queryResult().terminatedEarly()); - if (trackTotalHits == -1) { - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(0L)); - } else { - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(7L)); - } - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(7)); - assertProfileData(context, "BooleanQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), greaterThanOrEqualTo(7L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - assertThat(query.getProfiledChildren(), hasSize(2)); - assertThat(query.getProfiledChildren().get(0).getQueryName(), equalTo("TermQuery")); - assertThat(query.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("score_count"), greaterThan(0L)); - - assertThat(query.getProfiledChildren().get(1).getQueryName(), equalTo("TermQuery")); - assertThat(query.getProfiledChildren().get(1).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("score_count"), greaterThan(0L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_terminate_after_count")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_top_hits")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }); - } - - reader.close(); - dir.close(); - } - - public void testIndexSortingEarlyTermination() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("rank", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; ++i) { - Document doc = new Document(); - if (randomBoolean()) { - doc.add(new StringField("foo", "bar", Store.NO)); - } - if (randomBoolean()) { - doc.add(new StringField("foo", "baz", Store.NO)); - } - doc.add(new NumericDocValuesField("rank", numDocs - i)); - w.addDocument(doc); - } - w.close(); - - final IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - context.setSize(1); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.sort(new SortAndFormats(sort, new DocValueFormat[] { DocValueFormat.RAW })); - - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0], instanceOf(FieldDoc.class)); - FieldDoc fieldDoc = (FieldDoc) context.queryResult().topDocs().topDocs.scoreDocs[0]; - assertThat(fieldDoc.fields[0], equalTo(1)); - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - { - context.parsedPostFilter(new ParsedQuery(new MinDocQuery(1))); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(numDocs - 1L)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0], instanceOf(FieldDoc.class)); - assertThat(fieldDoc.fields[0], anyOf(equalTo(1), equalTo(2))); - assertProfileData(context, collector -> { - assertThat(collector.getReason(), equalTo("search_post_filter")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_top_hits")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }, (query) -> { - assertThat(query.getQueryName(), equalTo("MinDocQuery")); - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, (query) -> { - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertThat(query.getQueryName(), equalTo("ConstantScoreQuery")); - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }); - context.parsedPostFilter(null); - } - - { - context.setSearcher(newEarlyTerminationContextSearcher(reader, 1, executor)); - context.trackTotalHitsUpTo(SearchContext.TRACK_TOTAL_HITS_DISABLED); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0], instanceOf(FieldDoc.class)); - assertThat(fieldDoc.fields[0], anyOf(equalTo(1), equalTo(2))); - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0], instanceOf(FieldDoc.class)); - assertThat(fieldDoc.fields[0], anyOf(equalTo(1), equalTo(2))); - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - } - - reader.close(); - dir.close(); - } - - public void testIndexSortScrollOptimization() throws Exception { - Directory dir = newDirectory(); - final Sort indexSort = new Sort(new SortField("rank", SortField.Type.INT), new SortField("tiebreaker", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(indexSort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; ++i) { - Document doc = new Document(); - doc.add(new NumericDocValuesField("rank", random().nextInt())); - doc.add(new NumericDocValuesField("tiebreaker", i)); - w.addDocument(doc); - } - if (randomBoolean()) { - w.forceMerge(randomIntBetween(1, 10)); - } - w.close(); - - final IndexReader reader = DirectoryReader.open(dir); - List searchSortAndFormats = new ArrayList<>(); - searchSortAndFormats.add(new SortAndFormats(indexSort, new DocValueFormat[] { DocValueFormat.RAW, DocValueFormat.RAW })); - // search sort is a prefix of the index sort - searchSortAndFormats.add(new SortAndFormats(new Sort(indexSort.getSort()[0]), new DocValueFormat[] { DocValueFormat.RAW })); - for (SortAndFormats searchSortAndFormat : searchSortAndFormats) { - ScrollContext scrollContext = new ScrollContext(); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor), scrollContext); - context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); - scrollContext.lastEmittedDoc = null; - scrollContext.maxScore = Float.NaN; - scrollContext.totalHits = null; - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.setSize(10); - context.sort(searchSortAndFormat); - - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.terminateAfter(), equalTo(0)); - assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - int sizeMinus1 = context.queryResult().topDocs().topDocs.scoreDocs.length - 1; - FieldDoc lastDoc = (FieldDoc) context.queryResult().topDocs().topDocs.scoreDocs[sizeMinus1]; - - context.setSearcher(newEarlyTerminationContextSearcher(reader, 10, executor)); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertNull(context.queryResult().terminatedEarly()); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.terminateAfter(), equalTo(0)); - assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - assertThat(query.getProfiledChildren(), hasSize(1)); - assertThat(query.getProfiledChildren().get(0).getQueryName(), equalTo("SearchAfterSortedDocQuery")); - assertThat(query.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - FieldDoc firstDoc = (FieldDoc) context.queryResult().topDocs().topDocs.scoreDocs[0]; - for (int i = 0; i < searchSortAndFormat.sort.getSort().length; i++) { - @SuppressWarnings("unchecked") - FieldComparator comparator = (FieldComparator) searchSortAndFormat.sort.getSort()[i].getComparator(i, true); - int cmp = comparator.compareValues(firstDoc.fields[i], lastDoc.fields[i]); - if (cmp == 0) { - continue; - } - assertThat(cmp, equalTo(1)); - break; - } - } - reader.close(); - dir.close(); - } - - public void testDisableTopScoreCollection() throws Exception { - Directory dir = newDirectory(); - IndexWriterConfig iwc = newIndexWriterConfig(new StandardAnalyzer()); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - Document doc = new Document(); - final int numDocs = 2 * scaledRandomIntBetween(50, 450); - for (int i = 0; i < numDocs; i++) { - doc.clear(); - if (i % 2 == 0) { - doc.add(new TextField("title", "foo bar", Store.NO)); - } else { - doc.add(new TextField("title", "foo", Store.NO)); - } - w.addDocument(doc); - } - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - Query q = new SpanNearQuery.Builder("title", true).addClause(new SpanTermQuery(new Term("title", "foo"))) - .addClause(new SpanTermQuery(new Term("title", "bar"))) - .build(); - - context.parsedQuery(new ParsedQuery(q)); - context.setSize(3); - context.trackTotalHitsUpTo(3); - TopDocsCollectorContext topDocsContext = TopDocsCollectorContext.createTopDocsCollectorContext(context, false); - assertEquals(topDocsContext.create(null).scoreMode(), org.apache.lucene.search.ScoreMode.COMPLETE); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertEquals(numDocs / 2, context.queryResult().topDocs().topDocs.totalHits.value); - assertEquals(context.queryResult().topDocs().topDocs.totalHits.relation, TotalHits.Relation.EQUAL_TO); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(3)); - assertProfileData(context, "SpanNearQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - 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); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertEquals(numDocs / 2, context.queryResult().topDocs().topDocs.totalHits.value); - assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(3)); - assertEquals(context.queryResult().topDocs().topDocs.totalHits.relation, TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO); - // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery - // see: https://github.com/apache/lucene/pull/672 - assertProfileData(context, "ConstantScoreQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(0L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - reader.close(); - dir.close(); - } - - public void testMinScore() throws Exception { - Directory dir = newDirectory(); - IndexWriterConfig iwc = newIndexWriterConfig(); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - for (int i = 0; i < 10; i++) { - Document doc = new Document(); - doc.add(new StringField("foo", "bar", Store.NO)); - doc.add(new StringField("filter", "f1", Store.NO)); - w.addDocument(doc); - } - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.parsedQuery( - new ParsedQuery( - new BooleanQuery.Builder().add(new TermQuery(new Term("foo", "bar")), Occur.MUST) - .add(new TermQuery(new Term("filter", "f1")), Occur.SHOULD) - .build() - ) - ); - context.minimumScore(0.01f); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.setSize(1); - context.trackTotalHitsUpTo(5); - - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertEquals(10, context.queryResult().topDocs().topDocs.totalHits.value); - assertProfileData(context, "BooleanQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), equalTo(10L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - assertThat(query.getProfiledChildren(), hasSize(2)); - assertThat(query.getProfiledChildren().get(0).getQueryName(), equalTo("TermQuery")); - assertThat(query.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - assertThat(query.getProfiledChildren().get(1).getQueryName(), equalTo("TermQuery")); - assertThat(query.getProfiledChildren().get(1).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_min_score")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), hasSize(1)); - assertThat(collector.getProfiledChildren().get(0).getReason(), equalTo("search_top_hits")); - assertThat(collector.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - }); - - reader.close(); - dir.close(); - } - - public void testMaxScore() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("filter", SortField.Type.STRING)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - - final int numDocs = scaledRandomIntBetween(600, 900); - for (int i = 0; i < numDocs; i++) { - Document doc = new Document(); - doc.add(new StringField("foo", "bar", Store.NO)); - doc.add(new StringField("filter", "f1" + ((i > 0) ? " " + Integer.toString(i) : ""), Store.NO)); - doc.add(new SortedDocValuesField("filter", newBytesRef("f1" + ((i > 0) ? " " + Integer.toString(i) : "")))); - w.addDocument(doc); - } - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); - context.trackScores(true); - context.parsedQuery( - new ParsedQuery( - new BooleanQuery.Builder().add(new TermQuery(new Term("foo", "bar")), Occur.MUST) - .add(new TermQuery(new Term("filter", "f1")), Occur.SHOULD) - .build() - ) - ); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.setSize(1); - context.trackTotalHitsUpTo(5); - - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertFalse(Float.isNaN(context.queryResult().getMaxScore())); - assertEquals(1, context.queryResult().topDocs().topDocs.scoreDocs.length); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, greaterThanOrEqualTo(6L)); - assertProfileData(context, "BooleanQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), greaterThanOrEqualTo(6L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - assertThat(query.getProfiledChildren(), hasSize(2)); - assertThat(query.getProfiledChildren().get(0).getQueryName(), equalTo("TermQuery")); - assertThat(query.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - assertThat(query.getProfiledChildren().get(1).getQueryName(), equalTo("TermQuery")); - assertThat(query.getProfiledChildren().get(1).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - context.sort(new SortAndFormats(sort, new DocValueFormat[] { DocValueFormat.RAW })); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertFalse(Float.isNaN(context.queryResult().getMaxScore())); - assertEquals(1, context.queryResult().topDocs().topDocs.scoreDocs.length); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, greaterThanOrEqualTo(6L)); - assertProfileData(context, "BooleanQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), greaterThanOrEqualTo(6L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - assertThat(query.getProfiledChildren(), hasSize(2)); - assertThat(query.getProfiledChildren().get(0).getQueryName(), equalTo("TermQuery")); - assertThat(query.getProfiledChildren().get(0).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(0).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - - assertThat(query.getProfiledChildren().get(1).getQueryName(), equalTo("TermQuery")); - assertThat(query.getProfiledChildren().get(1).getTime(), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getProfiledChildren().get(1).getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - reader.close(); - dir.close(); - } - - public void testCollapseQuerySearchResults() throws Exception { - Directory dir = newDirectory(); - final Sort sort = new Sort(new SortField("user", SortField.Type.INT)); - IndexWriterConfig iwc = newIndexWriterConfig().setIndexSort(sort); - RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc); - - // Always end up with uneven buckets so collapsing is predictable - final int numDocs = 2 * scaledRandomIntBetween(600, 900) - 1; - for (int i = 0; i < numDocs; i++) { - Document doc = new Document(); - doc.add(new StringField("foo", "bar", Store.NO)); - doc.add(new NumericDocValuesField("user", i & 1)); - w.addDocument(doc); - } - w.close(); - - IndexReader reader = DirectoryReader.open(dir); - QueryShardContext queryShardContext = mock(QueryShardContext.class); - when(queryShardContext.fieldMapper("user")).thenReturn( - new NumberFieldType("user", NumberType.INTEGER, true, false, true, false, null, Collections.emptyMap()) - ); - - TestSearchContext context = new TestSearchContext(queryShardContext, indexShard, newContextSearcher(reader, executor)); - context.collapse(new CollapseBuilder("user").build(context.getQueryShardContext())); - context.trackScores(true); - context.parsedQuery(new ParsedQuery(new TermQuery(new Term("foo", "bar")))); - context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); - context.setSize(2); - context.trackTotalHitsUpTo(5); - - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertFalse(Float.isNaN(context.queryResult().getMaxScore())); - assertEquals(2, context.queryResult().topDocs().topDocs.scoreDocs.length); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs, instanceOf(CollapseTopFieldDocs.class)); - - assertProfileData(context, "TermQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), greaterThanOrEqualTo(6L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - assertThat(query.getProfiledChildren(), empty()); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - context.sort(new SortAndFormats(sort, new DocValueFormat[] { DocValueFormat.RAW })); - QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers(), queryPhaseSearcher); - assertFalse(Float.isNaN(context.queryResult().getMaxScore())); - assertEquals(2, context.queryResult().topDocs().topDocs.scoreDocs.length); - assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); - assertThat(context.queryResult().topDocs().topDocs, instanceOf(CollapseTopFieldDocs.class)); - - assertProfileData(context, "TermQuery", query -> { - assertThat(query.getTimeBreakdown().keySet(), not(empty())); - assertThat(query.getTimeBreakdown().get("score"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("score_count"), greaterThanOrEqualTo(6L)); - assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); - assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); - assertThat(query.getProfiledChildren(), empty()); - }, collector -> { - assertThat(collector.getReason(), equalTo("search_top_hits")); - assertThat(collector.getTime(), greaterThan(0L)); - assertThat(collector.getProfiledChildren(), empty()); - }); - - reader.close(); - dir.close(); - } - - private void assertProfileData(SearchContext context, String type, Consumer query, Consumer collector) - throws IOException { - assertProfileData(context, collector, (profileResult) -> { - assertThat(profileResult.getQueryName(), equalTo(type)); - assertThat(profileResult.getTime(), greaterThan(0L)); - query.accept(profileResult); - }); - } - - private void assertProfileData(SearchContext context, Consumer collector, Consumer query1) - throws IOException { - assertProfileData(context, Arrays.asList(query1), collector, false); - } - - private void assertProfileData( - SearchContext context, - Consumer collector, - Consumer query1, - Consumer query2 - ) throws IOException { - assertProfileData(context, Arrays.asList(query1, query2), collector, false); - } - - private final void assertProfileData( - SearchContext context, - List> queries, - Consumer collector, - boolean debug - ) throws IOException { - assertThat(context.getProfilers(), not(nullValue())); - - final ProfileShardResult result = SearchProfileShardResults.buildShardResults(context.getProfilers(), null); - if (debug) { - final SearchProfileShardResults results = new SearchProfileShardResults( - Collections.singletonMap(indexShard.shardId().toString(), result) - ); - - try (final XContentBuilder builder = JsonXContent.contentBuilder().prettyPrint()) { - builder.startObject(); - results.toXContent(builder, ToXContent.EMPTY_PARAMS); - builder.endObject(); - builder.flush(); - - final OutputStream out = builder.getOutputStream(); - assertThat(out, instanceOf(ByteArrayOutputStream.class)); - - logger.info(new String(((ByteArrayOutputStream) out).toByteArray(), StandardCharsets.UTF_8)); - } - } - - assertThat(result.getQueryProfileResults(), hasSize(1)); - - final QueryProfileShardResult queryProfileShardResult = result.getQueryProfileResults().get(0); - assertThat(queryProfileShardResult.getQueryResults(), hasSize(queries.size())); - - for (int i = 0; i < queries.size(); ++i) { - queries.get(i).accept(queryProfileShardResult.getQueryResults().get(i)); - } - - collector.accept(queryProfileShardResult.getCollectorResult()); - } - - private static ContextIndexSearcher newContextSearcher(IndexReader reader, ExecutorService executor) throws IOException { - return new ContextIndexSearcher( - reader, - IndexSearcher.getDefaultSimilarity(), - IndexSearcher.getDefaultQueryCache(), - IndexSearcher.getDefaultQueryCachingPolicy(), - true, - executor - ); - } - - private static ContextIndexSearcher newEarlyTerminationContextSearcher(IndexReader reader, int size, ExecutorService executor) - throws IOException { - return new ContextIndexSearcher( - reader, - IndexSearcher.getDefaultSimilarity(), - IndexSearcher.getDefaultQueryCache(), - IndexSearcher.getDefaultQueryCachingPolicy(), - true, - executor - ) { - - @Override - public void search(List leaves, Weight weight, Collector collector) throws IOException { - final Collector in = new AssertingEarlyTerminationFilterCollector(collector, size); - super.search(leaves, weight, in); - } - }; - } - - private static class AssertingEarlyTerminationFilterCollector extends FilterCollector { - private final int size; - - AssertingEarlyTerminationFilterCollector(Collector in, int size) { - super(in); - this.size = size; - } - - @Override - public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException { - final LeafCollector in = super.getLeafCollector(context); - return new FilterLeafCollector(in) { - int collected; - - @Override - public void collect(int doc) throws IOException { - assert collected <= size : "should not collect more than " + size + " doc per segment, got " + collected; - ++collected; - super.collect(doc); - } - }; - } - } -} diff --git a/server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java b/server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java index 0f819ad02bf0e..bd94e414cb903 100644 --- a/server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java @@ -37,7 +37,8 @@ protected FeatureFlagSettings( FeatureFlags.SEGMENT_REPLICATION_EXPERIMENTAL_SETTING, FeatureFlags.REMOTE_STORE_SETTING, FeatureFlags.EXTENSIONS_SETTING, - FeatureFlags.SEARCH_PIPELINE_SETTING + FeatureFlags.SEARCH_PIPELINE_SETTING, + FeatureFlags.CONCURRENT_SEGMENT_SEARCH_SETTING ) ) ); diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java index 34d0feeaeb525..91f3cc665ab58 100644 --- a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java +++ b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java @@ -52,6 +52,12 @@ public class FeatureFlags { */ public static final String SEARCH_PIPELINE = "opensearch.experimental.feature.search_pipeline.enabled"; + /** + * Gates the functionality of concurrently searching the segments + * Once the feature is ready for release, this feature flag can be removed. + */ + public static final String CONCURRENT_SEGMENT_SEARCH = "opensearch.experimental.feature.concurrent_segment_search.enabled"; + /** * Should store the settings from opensearch.yml. */ @@ -91,4 +97,10 @@ public static boolean isEnabled(String featureFlagName) { public static final Setting EXTENSIONS_SETTING = Setting.boolSetting(EXTENSIONS, false, Property.NodeScope); public static final Setting SEARCH_PIPELINE_SETTING = Setting.boolSetting(SEARCH_PIPELINE, false, Property.NodeScope); + + public static final Setting CONCURRENT_SEGMENT_SEARCH_SETTING = Setting.boolSetting( + CONCURRENT_SEGMENT_SEARCH, + false, + Property.NodeScope + ); } diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java index a3364cd1c508c..74fa3dcddbb7d 100644 --- a/server/src/main/java/org/opensearch/node/Node.java +++ b/server/src/main/java/org/opensearch/node/Node.java @@ -237,7 +237,6 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.CountDownLatch; -import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; @@ -1014,8 +1013,7 @@ protected Node( searchModule.getQueryPhase(), searchModule.getFetchPhase(), responseCollectorService, - circuitBreakerService, - searchModule.getIndexSearcherExecutor(threadPool) + circuitBreakerService ); final List> tasksExecutors = pluginsService.filterPlugins(PersistentTaskPlugin.class) @@ -1611,8 +1609,7 @@ protected SearchService newSearchService( QueryPhase queryPhase, FetchPhase fetchPhase, ResponseCollectorService responseCollectorService, - CircuitBreakerService circuitBreakerService, - Executor indexSearcherExecutor + CircuitBreakerService circuitBreakerService ) { return new SearchService( clusterService, @@ -1623,8 +1620,7 @@ protected SearchService newSearchService( queryPhase, fetchPhase, responseCollectorService, - circuitBreakerService, - indexSearcherExecutor + circuitBreakerService ); } diff --git a/server/src/main/java/org/opensearch/search/SearchModule.java b/server/src/main/java/org/opensearch/search/SearchModule.java index 30614b6eea0fa..addce2ee1d961 100644 --- a/server/src/main/java/org/opensearch/search/SearchModule.java +++ b/server/src/main/java/org/opensearch/search/SearchModule.java @@ -34,7 +34,6 @@ import org.apache.lucene.search.BooleanQuery; import org.opensearch.common.NamedRegistry; -import org.opensearch.common.Nullable; import org.opensearch.core.ParseField; import org.opensearch.common.geo.GeoShapeType; import org.opensearch.common.geo.ShapesAvailability; @@ -277,14 +276,12 @@ import org.opensearch.search.suggest.phrase.StupidBackoff; import org.opensearch.search.suggest.term.TermSuggestion; import org.opensearch.search.suggest.term.TermSuggestionBuilder; -import org.opensearch.threadpool.ThreadPool; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.concurrent.ExecutorService; import java.util.function.Consumer; import java.util.function.Function; @@ -319,7 +316,6 @@ public class SearchModule { private final List namedXContents = new ArrayList<>(); private final ValuesSourceRegistry valuesSourceRegistry; private final QueryPhaseSearcher queryPhaseSearcher; - private final SearchPlugin.ExecutorServiceProvider indexSearcherExecutorProvider; /** * Constructs a new SearchModule object @@ -347,7 +343,6 @@ public SearchModule(Settings settings, List plugins) { registerShapes(); registerIntervalsSourceProviders(); queryPhaseSearcher = registerQueryPhaseSearcher(plugins); - indexSearcherExecutorProvider = registerIndexSearcherExecutorProvider(plugins); namedWriteables.addAll(SortValue.namedWriteables()); } @@ -1255,24 +1250,6 @@ private QueryPhaseSearcher registerQueryPhaseSearcher(List plugins return searcher; } - private SearchPlugin.ExecutorServiceProvider registerIndexSearcherExecutorProvider(List plugins) { - SearchPlugin.ExecutorServiceProvider provider = null; - - for (SearchPlugin plugin : plugins) { - final Optional providerOpt = plugin.getIndexSearcherExecutorProvider(); - - if (provider == null) { - provider = providerOpt.orElse(null); - } else if (providerOpt.isPresent()) { - throw new IllegalStateException( - "The index searcher executor is already assigned but more than one are provided by the plugins" - ); - } - } - - return provider; - } - public FetchPhase getFetchPhase() { return new FetchPhase(fetchSubPhases); } @@ -1280,8 +1257,4 @@ public FetchPhase getFetchPhase() { public QueryPhase getQueryPhase() { return (queryPhaseSearcher == null) ? new QueryPhase() : new QueryPhase(queryPhaseSearcher); } - - public @Nullable ExecutorService getIndexSearcherExecutor(ThreadPool pool) { - return (indexSearcherExecutorProvider == null) ? null : indexSearcherExecutorProvider.getExecutor(pool); - } } diff --git a/server/src/main/java/org/opensearch/search/SearchService.java b/server/src/main/java/org/opensearch/search/SearchService.java index cc1ee60dbb6bf..1a5c5fcb80747 100644 --- a/server/src/main/java/org/opensearch/search/SearchService.java +++ b/server/src/main/java/org/opensearch/search/SearchService.java @@ -67,6 +67,7 @@ import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.BigArrays; import org.opensearch.common.util.CollectionUtils; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.common.util.concurrent.ConcurrentCollections; import org.opensearch.common.util.concurrent.ConcurrentMapLong; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; @@ -306,8 +307,7 @@ public SearchService( QueryPhase queryPhase, FetchPhase fetchPhase, ResponseCollectorService responseCollectorService, - CircuitBreakerService circuitBreakerService, - Executor indexSearcherExecutor + CircuitBreakerService circuitBreakerService ) { Settings settings = clusterService.getSettings(); this.threadPool = threadPool; @@ -323,7 +323,9 @@ public SearchService( settings, circuitBreakerService.getBreaker(CircuitBreaker.REQUEST) ); - this.indexSearcherExecutor = indexSearcherExecutor; + this.indexSearcherExecutor = FeatureFlags.isEnabled(FeatureFlags.CONCURRENT_SEGMENT_SEARCH) + ? threadPool.executor(ThreadPool.Names.INDEX_SEARCHER) + : null; TimeValue keepAliveInterval = KEEPALIVE_INTERVAL_SETTING.get(settings); setKeepAlives(DEFAULT_KEEPALIVE_SETTING.get(settings), MAX_KEEPALIVE_SETTING.get(settings)); diff --git a/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java b/server/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java similarity index 96% rename from sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java rename to server/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java index 65f339838a40b..7a94cec4b31e2 100644 --- a/sandbox/plugins/concurrent-search/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java +++ b/server/src/main/java/org/opensearch/search/query/ConcurrentQueryPhaseSearcher.java @@ -18,6 +18,7 @@ import org.apache.lucene.search.Collector; import org.apache.lucene.search.CollectorManager; import org.apache.lucene.search.Query; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.search.internal.ContextIndexSearcher; import org.opensearch.search.internal.SearchContext; import org.opensearch.search.profile.query.ProfileCollectorManager; @@ -113,7 +114,7 @@ private static boolean searchWithCollectorManager( } private static boolean allowConcurrentSegmentSearch(final ContextIndexSearcher searcher) { - return (searcher.getExecutor() != null); + return (searcher.getExecutor() != null) && FeatureFlags.isEnabled(FeatureFlags.CONCURRENT_SEGMENT_SEARCH); } } diff --git a/server/src/main/java/org/opensearch/search/query/QueryPhase.java b/server/src/main/java/org/opensearch/search/query/QueryPhase.java index 405efeb37cb70..c8b54e8c7bf6d 100644 --- a/server/src/main/java/org/opensearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/opensearch/search/query/QueryPhase.java @@ -36,6 +36,7 @@ import org.apache.logging.log4j.Logger; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.LeafReaderContext; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.lucene.queries.SearchAfterSortedDocQuery; import org.apache.lucene.search.BooleanClause; import org.apache.lucene.search.BooleanQuery; @@ -88,7 +89,9 @@ public class QueryPhase { private static final Logger LOGGER = LogManager.getLogger(QueryPhase.class); // TODO: remove this property public static final boolean SYS_PROP_REWRITE_SORT = Booleans.parseBoolean(System.getProperty("opensearch.search.rewrite_sort", "true")); - public static final QueryPhaseSearcher DEFAULT_QUERY_PHASE_SEARCHER = new DefaultQueryPhaseSearcher(); + public static final QueryPhaseSearcher DEFAULT_QUERY_PHASE_SEARCHER = FeatureFlags.isEnabled(FeatureFlags.CONCURRENT_SEGMENT_SEARCH) + ? new ConcurrentQueryPhaseSearcher() + : new DefaultQueryPhaseSearcher(); private final QueryPhaseSearcher queryPhaseSearcher; private final AggregationPhase aggregationPhase; diff --git a/server/src/main/java/org/opensearch/threadpool/ThreadPool.java b/server/src/main/java/org/opensearch/threadpool/ThreadPool.java index 104af6945dfe8..acb6defb558c4 100644 --- a/server/src/main/java/org/opensearch/threadpool/ThreadPool.java +++ b/server/src/main/java/org/opensearch/threadpool/ThreadPool.java @@ -111,6 +111,7 @@ public static class Names { public static final String TRANSLOG_TRANSFER = "translog_transfer"; public static final String TRANSLOG_SYNC = "translog_sync"; public static final String REMOTE_PURGE = "remote_purge"; + public static final String INDEX_SEARCHER = "index_searcher"; } /** @@ -178,6 +179,7 @@ public static ThreadPoolType fromType(String type) { map.put(Names.TRANSLOG_TRANSFER, ThreadPoolType.SCALING); map.put(Names.TRANSLOG_SYNC, ThreadPoolType.FIXED); map.put(Names.REMOTE_PURGE, ThreadPoolType.SCALING); + map.put(Names.INDEX_SEARCHER, ThreadPoolType.FIXED); THREAD_POOL_TYPES = Collections.unmodifiableMap(map); } @@ -256,6 +258,7 @@ public ThreadPool( ); builders.put(Names.TRANSLOG_SYNC, new FixedExecutorBuilder(settings, Names.TRANSLOG_SYNC, allocatedProcessors * 4, 10000)); builders.put(Names.REMOTE_PURGE, new ScalingExecutorBuilder(Names.REMOTE_PURGE, 1, halfProcMaxAt5, TimeValue.timeValueMinutes(5))); + builders.put(Names.INDEX_SEARCHER, new FixedExecutorBuilder(settings, Names.INDEX_SEARCHER, allocatedProcessors, 1000, false)); for (final ExecutorBuilder builder : customBuilders) { if (builders.containsKey(builder.name())) { diff --git a/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java b/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java index da2a70514a716..8ad391330e690 100644 --- a/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java +++ b/server/src/test/java/org/opensearch/search/profile/query/QueryProfilerTests.java @@ -58,16 +58,21 @@ import org.apache.lucene.search.Weight; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.util.TestUtil; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.common.util.io.IOUtils; import org.opensearch.search.internal.ContextIndexSearcher; import org.opensearch.search.profile.ProfileResult; import org.opensearch.test.OpenSearchTestCase; import org.junit.After; import org.junit.Before; +import org.opensearch.threadpool.ThreadPool; import java.io.IOException; import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.greaterThan; @@ -77,6 +82,7 @@ public class QueryProfilerTests extends OpenSearchTestCase { private Directory dir; private IndexReader reader; private ContextIndexSearcher searcher; + private ExecutorService executor; @Before public void setUp() throws Exception { @@ -96,13 +102,14 @@ public void setUp() throws Exception { } reader = w.getReader(); w.close(); + executor = (FeatureFlags.isEnabled(FeatureFlags.CONCURRENT_SEGMENT_SEARCH) ? Executors.newFixedThreadPool(5) : null); searcher = new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), ALWAYS_CACHE_POLICY, true, - null + executor ); } @@ -116,6 +123,10 @@ public void tearDown() throws Exception { assertThat(cache.getTotalCount(), equalTo(cache.getMissCount())); assertThat(cache.getCacheSize(), equalTo(0L)); + if (executor != null) { + ThreadPool.terminate(executor, 10, TimeUnit.SECONDS); + } + IOUtils.close(reader, dir); dir = null; reader = null; @@ -123,7 +134,7 @@ public void tearDown() throws Exception { } public void testBasic() throws IOException { - QueryProfiler profiler = new QueryProfiler(false); + QueryProfiler profiler = new QueryProfiler(executor != null); searcher.setProfiler(profiler); Query query = new TermQuery(new Term("foo", "bar")); searcher.search(query, 1); @@ -149,7 +160,7 @@ public void testBasic() throws IOException { } public void testNoScoring() throws IOException { - QueryProfiler profiler = new QueryProfiler(false); + QueryProfiler profiler = new QueryProfiler(executor != null); searcher.setProfiler(profiler); Query query = new TermQuery(new Term("foo", "bar")); searcher.search(query, 1, Sort.INDEXORDER); // scores are not needed @@ -175,7 +186,7 @@ public void testNoScoring() throws IOException { } public void testUseIndexStats() throws IOException { - QueryProfiler profiler = new QueryProfiler(false); + QueryProfiler profiler = new QueryProfiler(executor != null); searcher.setProfiler(profiler); Query query = new TermQuery(new Term("foo", "bar")); searcher.count(query); // will use index stats @@ -189,7 +200,7 @@ public void testUseIndexStats() throws IOException { } public void testApproximations() throws IOException { - QueryProfiler profiler = new QueryProfiler(false); + QueryProfiler profiler = new QueryProfiler(executor != null); searcher.setProfiler(profiler); Query query = new RandomApproximationQuery(new TermQuery(new Term("foo", "bar")), random()); searcher.count(query); diff --git a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java index 18225ba887416..c4be5fcb0ecbc 100644 --- a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java +++ b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java @@ -86,6 +86,7 @@ import org.opensearch.action.search.SearchShardTask; import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.index.mapper.DateFieldMapper; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; @@ -110,8 +111,12 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import static org.hamcrest.Matchers.anyOf; import static org.hamcrest.Matchers.equalTo; @@ -133,6 +138,8 @@ public class QueryPhaseTests extends IndexShardTestCase { private IndexShard indexShard; + private ExecutorService executor; + @Override public Settings threadPoolSettings() { return Settings.builder().put(super.threadPoolSettings()).put("thread_pool.search.min_queue_size", 10).build(); @@ -141,17 +148,23 @@ public Settings threadPoolSettings() { @Override public void setUp() throws Exception { super.setUp(); + this.executor = (FeatureFlags.isEnabled(FeatureFlags.CONCURRENT_SEGMENT_SEARCH)) ? Executors.newFixedThreadPool(5) : null; indexShard = newShard(true); } @Override public void tearDown() throws Exception { super.tearDown(); + if (executor != null) { + ThreadPool.terminate(executor, 10, TimeUnit.SECONDS); + } closeShards(indexShard); } private void countTestCase(Query query, IndexReader reader, boolean shouldCollectSearch, boolean shouldCollectCount) throws Exception { - ContextIndexSearcher searcher = shouldCollectSearch ? newContextSearcher(reader) : newEarlyTerminationContextSearcher(reader, 0); + ContextIndexSearcher searcher = shouldCollectSearch + ? newContextSearcher(reader, executor) + : newEarlyTerminationContextSearcher(reader, 0, executor); TestSearchContext context = new TestSearchContext(null, indexShard, searcher); context.parsedQuery(new ParsedQuery(query)); context.setSize(0); @@ -160,8 +173,8 @@ private void countTestCase(Query query, IndexReader reader, boolean shouldCollec assertFalse(rescore); ContextIndexSearcher countSearcher = shouldCollectCount - ? newContextSearcher(reader) - : newEarlyTerminationContextSearcher(reader, 0); + ? newContextSearcher(reader, executor) + : newEarlyTerminationContextSearcher(reader, 0, executor); assertEquals(countSearcher.count(query), context.queryResult().topDocs().topDocs.totalHits.value); } @@ -236,14 +249,14 @@ public void testPostFilterDisablesCountOptimization() throws Exception { IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0)); + TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0, executor)); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); QueryPhase.executeInternal(context.withCleanQueryResult()); assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value); - context.setSearcher(newContextSearcher(reader)); + context.setSearcher(newContextSearcher(reader, executor)); context.parsedPostFilter(new ParsedQuery(new MatchNoDocsQuery())); QueryPhase.executeInternal(context.withCleanQueryResult()); assertEquals(0, context.queryResult().topDocs().topDocs.totalHits.value); @@ -265,7 +278,7 @@ public void testTerminateAfterWithFilter() throws Exception { IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); @@ -291,7 +304,7 @@ public void testMinScoreDisablesCountOptimization() throws Exception { w.close(); IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0)); + TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0, executor)); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); context.setSize(0); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); @@ -316,7 +329,7 @@ public void testQueryCapturesThreadPoolStats() throws Exception { } w.close(); IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); @@ -340,7 +353,7 @@ public void testInOrderScrollOptimization() throws Exception { w.close(); IndexReader reader = DirectoryReader.open(dir); ScrollContext scrollContext = new ScrollContext(); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader), scrollContext); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor), scrollContext); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); context.sort(new SortAndFormats(sort, new DocValueFormat[] { DocValueFormat.RAW })); scrollContext.lastEmittedDoc = null; @@ -356,7 +369,7 @@ public void testInOrderScrollOptimization() throws Exception { assertThat(context.terminateAfter(), equalTo(0)); assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); - context.setSearcher(newEarlyTerminationContextSearcher(reader, size)); + context.setSearcher(newEarlyTerminationContextSearcher(reader, size, executor)); QueryPhase.executeInternal(context.withCleanQueryResult()); assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); @@ -383,14 +396,14 @@ public void testTerminateAfterEarlyTermination() throws Exception { } w.close(); final IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); context.terminateAfter(numDocs); { context.setSize(10); - final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(); + final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(executor); context.queryCollectorManagers().put(TotalHitCountCollector.class, manager); QueryPhase.executeInternal(context.withCleanQueryResult()); assertFalse(context.queryResult().terminatedEarly()); @@ -441,7 +454,7 @@ public void testTerminateAfterEarlyTermination() throws Exception { } { context.setSize(1); - final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(); + final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(executor, 1); context.queryCollectorManagers().put(TotalHitCountCollector.class, manager); QueryPhase.executeInternal(context.withCleanQueryResult()); assertTrue(context.queryResult().terminatedEarly()); @@ -452,7 +465,7 @@ public void testTerminateAfterEarlyTermination() throws Exception { } { context.setSize(0); - final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(); + final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(executor, 1); context.queryCollectorManagers().put(TotalHitCountCollector.class, manager); QueryPhase.executeInternal(context.withCleanQueryResult()); assertTrue(context.queryResult().terminatedEarly()); @@ -466,7 +479,7 @@ public void testTerminateAfterEarlyTermination() throws Exception { context.setSize(0); for (int trackTotalHits : new int[] { -1, 3, 76, 100 }) { context.trackTotalHitsUpTo(trackTotalHits); - final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(); + final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(executor); context.queryCollectorManagers().put(TotalHitCountCollector.class, manager); QueryPhase.executeInternal(context.withCleanQueryResult()); assertTrue(context.queryResult().terminatedEarly()); @@ -476,7 +489,13 @@ public void testTerminateAfterEarlyTermination() throws Exception { assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) Math.min(trackTotalHits, 10))); } assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0)); - assertThat(manager.getTotalHits(), equalTo(10)); + // The concurrent search terminates the collection when the number of hits is reached by each + // concurrent collector. In this case, in general, the number of results are multiplied by the number of + // slices (as the unit of concurrency). To address that, we have to use the shared global state, + // much as HitsThresholdChecker does. + if (executor == null) { + assertThat(manager.getTotalHits(), equalTo(10)); + } } context.terminateAfter(7); @@ -516,7 +535,7 @@ public void testIndexSortingEarlyTermination() throws Exception { w.close(); final IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); context.setSize(1); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); @@ -539,7 +558,7 @@ public void testIndexSortingEarlyTermination() throws Exception { assertThat(fieldDoc.fields[0], anyOf(equalTo(1), equalTo(2))); context.parsedPostFilter(null); - final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(sort); + final TestTotalHitCountCollectorManager manager = TestTotalHitCountCollectorManager.create(executor, sort); context.queryCollectorManagers().put(TotalHitCountCollector.class, manager); QueryPhase.executeInternal(context.withCleanQueryResult()); assertNull(context.queryResult().terminatedEarly()); @@ -547,12 +566,14 @@ public void testIndexSortingEarlyTermination() throws Exception { assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1)); assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0], instanceOf(FieldDoc.class)); assertThat(fieldDoc.fields[0], anyOf(equalTo(1), equalTo(2))); - assertThat(manager.getTotalHits(), equalTo(numDocs)); + // When searching concurrently, each executors short-circuits when "size" is reached, + // including total hits collector + assertThat(manager.getTotalHits(), lessThanOrEqualTo(numDocs)); context.queryCollectorManagers().clear(); } { - context.setSearcher(newEarlyTerminationContextSearcher(reader, 1)); + context.setSearcher(newEarlyTerminationContextSearcher(reader, 1, executor)); context.trackTotalHitsUpTo(SearchContext.TRACK_TOTAL_HITS_DISABLED); QueryPhase.executeInternal(context.withCleanQueryResult()); assertNull(context.queryResult().terminatedEarly()); @@ -594,7 +615,7 @@ public void testIndexSortScrollOptimization() throws Exception { searchSortAndFormats.add(new SortAndFormats(new Sort(indexSort.getSort()[0]), new DocValueFormat[] { DocValueFormat.RAW })); for (SortAndFormats searchSortAndFormat : searchSortAndFormats) { ScrollContext scrollContext = new ScrollContext(); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader), scrollContext); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor), scrollContext); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); scrollContext.lastEmittedDoc = null; scrollContext.maxScore = Float.NaN; @@ -611,7 +632,7 @@ public void testIndexSortScrollOptimization() throws Exception { int sizeMinus1 = context.queryResult().topDocs().topDocs.scoreDocs.length - 1; FieldDoc lastDoc = (FieldDoc) context.queryResult().topDocs().topDocs.scoreDocs[sizeMinus1]; - context.setSearcher(newEarlyTerminationContextSearcher(reader, 10)); + context.setSearcher(newEarlyTerminationContextSearcher(reader, 10, executor)); QueryPhase.executeInternal(context.withCleanQueryResult()); assertNull(context.queryResult().terminatedEarly()); assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); @@ -654,7 +675,7 @@ public void testDisableTopScoreCollection() throws Exception { w.close(); IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); Query q = new SpanNearQuery.Builder("title", true).addClause(new SpanTermQuery(new Term("title", "foo"))) .addClause(new SpanTermQuery(new Term("title", "bar"))) @@ -730,7 +751,7 @@ public void testEnhanceSortOnNumeric() throws Exception { // 1. Test a sort on long field { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader))); + TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); when(searchContext.mapperService()).thenReturn(mapperService); searchContext.sort(longSortAndFormats); searchContext.parsedQuery(query); @@ -742,7 +763,7 @@ public void testEnhanceSortOnNumeric() throws Exception { // 2. Test a sort on long field + date field { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader))); + TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); when(searchContext.mapperService()).thenReturn(mapperService); searchContext.sort(longDateSortAndFormats); searchContext.parsedQuery(query); @@ -754,7 +775,7 @@ public void testEnhanceSortOnNumeric() throws Exception { // 3. Test a sort on date field { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader))); + TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); when(searchContext.mapperService()).thenReturn(mapperService); searchContext.sort(dateSortAndFormats); searchContext.parsedQuery(query); @@ -766,7 +787,7 @@ public void testEnhanceSortOnNumeric() throws Exception { // 4. Test a sort on date field + long field { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader))); + TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); when(searchContext.mapperService()).thenReturn(mapperService); searchContext.sort(dateLongSortAndFormats); searchContext.parsedQuery(query); @@ -778,7 +799,7 @@ public void testEnhanceSortOnNumeric() throws Exception { // 5. Test that sort optimization is run when from > 0 and size = 0 { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader))); + TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); when(searchContext.mapperService()).thenReturn(mapperService); searchContext.sort(longSortAndFormats); searchContext.parsedQuery(query); @@ -791,7 +812,7 @@ public void testEnhanceSortOnNumeric() throws Exception { // 6. Test that sort optimization works with from = 0 and size= 0 { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader))); + TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); when(searchContext.mapperService()).thenReturn(mapperService); searchContext.sort(longSortAndFormats); searchContext.parsedQuery(query); @@ -802,7 +823,7 @@ public void testEnhanceSortOnNumeric() throws Exception { // 7. Test that sort optimization works with search after { - TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader))); + TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader, executor))); when(searchContext.mapperService()).thenReturn(mapperService); int afterDocument = (int) randomLongBetween(0, 50); long afterValue = firstValue + afterDocument; @@ -919,7 +940,7 @@ public void testMinScore() throws Exception { w.close(); IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.parsedQuery( new ParsedQuery( new BooleanQuery.Builder().add(new TermQuery(new Term("foo", "bar")), Occur.MUST) @@ -956,7 +977,7 @@ public void testMaxScore() throws Exception { w.close(); IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.trackScores(true); context.parsedQuery( new ParsedQuery( @@ -1012,7 +1033,7 @@ public void testCollapseQuerySearchResults() throws Exception { new NumberFieldType("user", NumberType.INTEGER, true, false, true, false, null, Collections.emptyMap()) ); - TestSearchContext context = new TestSearchContext(queryShardContext, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(queryShardContext, indexShard, newContextSearcher(reader, executor)); context.collapse(new CollapseBuilder("user").build(context.getQueryShardContext())); context.trackScores(true); context.parsedQuery(new ParsedQuery(new TermQuery(new Term("foo", "bar")))); @@ -1075,7 +1096,11 @@ public void testCancellationDuringPreprocess() throws IOException { w.close(); try (IndexReader reader = DirectoryReader.open(dir)) { - TestSearchContext context = new TestSearchContextWithRewriteAndCancellation(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContextWithRewriteAndCancellation( + null, + indexShard, + newContextSearcher(reader, executor) + ); PrefixQuery prefixQuery = new PrefixQuery(new Term("foo", "a"), MultiTermQuery.SCORING_BOOLEAN_REWRITE); context.parsedQuery(new ParsedQuery(prefixQuery)); SearchShardTask task = mock(SearchShardTask.class); @@ -1163,25 +1188,26 @@ public boolean lowLevelCancellation() { } } - private static ContextIndexSearcher newContextSearcher(IndexReader reader) throws IOException { + private static ContextIndexSearcher newContextSearcher(IndexReader reader, ExecutorService executor) throws IOException { return new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - null + executor ); } - private static ContextIndexSearcher newEarlyTerminationContextSearcher(IndexReader reader, int size) throws IOException { + private static ContextIndexSearcher newEarlyTerminationContextSearcher(IndexReader reader, int size, ExecutorService executor) + throws IOException { return new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - null + executor ) { @Override @@ -1193,28 +1219,62 @@ public void search(List leaves, Weight weight, Collector coll } private static class TestTotalHitCountCollectorManager extends TotalHitCountCollectorManager { + private int totalHits; private final TotalHitCountCollector collector; + private final Integer teminateAfter; - static TestTotalHitCountCollectorManager create() { - return create(null); + static TestTotalHitCountCollectorManager create(final ExecutorService executor) { + return create(executor, null, null); } - static TestTotalHitCountCollectorManager create(final Sort sort) { - return new TestTotalHitCountCollectorManager(new TotalHitCountCollector(), sort); + static TestTotalHitCountCollectorManager create(final ExecutorService executor, final Integer teminateAfter) { + return create(executor, null, teminateAfter); + } + + static TestTotalHitCountCollectorManager create(final ExecutorService executor, final Sort sort) { + return create(executor, sort, null); + } + + static TestTotalHitCountCollectorManager create(final ExecutorService executor, final Sort sort, final Integer teminateAfter) { + if (executor == null) { + return new TestTotalHitCountCollectorManager(new TotalHitCountCollector(), sort); + } else { + return new TestTotalHitCountCollectorManager(sort, teminateAfter); + } } private TestTotalHitCountCollectorManager(final TotalHitCountCollector collector, final Sort sort) { super(sort); this.collector = collector; + this.teminateAfter = null; + } + + private TestTotalHitCountCollectorManager(final Sort sort, final Integer teminateAfter) { + super(sort); + this.collector = null; + this.teminateAfter = teminateAfter; } @Override public TotalHitCountCollector newCollector() throws IOException { - return collector; + return (collector == null) ? super.newCollector() : collector; + } + + @Override + public ReduceableSearchResult reduce(Collection collectors) throws IOException { + final ReduceableSearchResult result = super.reduce(collectors); + totalHits = collectors.stream().mapToInt(TotalHitCountCollector::getTotalHits).sum(); + + if (teminateAfter != null) { + assertThat(totalHits, greaterThanOrEqualTo(teminateAfter)); + totalHits = Math.min(totalHits, teminateAfter); + } + + return result; } public int getTotalHits() { - return collector.getTotalHits(); + return (collector == null) ? totalHits : collector.getTotalHits(); } } diff --git a/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java index 98e23c90aa7ff..e45648fba3bd4 100644 --- a/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java +++ b/server/src/test/java/org/opensearch/search/query/QueryProfilePhaseTests.java @@ -44,6 +44,7 @@ import org.apache.lucene.tests.index.RandomIndexWriter; import org.opensearch.action.search.SearchShardTask; import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.common.xcontent.json.JsonXContent; @@ -66,6 +67,7 @@ import org.opensearch.search.profile.query.QueryProfileShardResult; import org.opensearch.search.sort.SortAndFormats; import org.opensearch.test.TestSearchContext; +import org.opensearch.threadpool.ThreadPool; import java.io.ByteArrayOutputStream; import java.io.IOException; @@ -75,6 +77,9 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import static org.hamcrest.CoreMatchers.not; @@ -93,6 +98,8 @@ public class QueryProfilePhaseTests extends IndexShardTestCase { private IndexShard indexShard; + private ExecutorService executor; + @Override public Settings threadPoolSettings() { return Settings.builder().put(super.threadPoolSettings()).put("thread_pool.search.min_queue_size", 10).build(); @@ -101,12 +108,16 @@ public Settings threadPoolSettings() { @Override public void setUp() throws Exception { super.setUp(); + this.executor = (FeatureFlags.isEnabled(FeatureFlags.CONCURRENT_SEGMENT_SEARCH)) ? Executors.newFixedThreadPool(5) : null; indexShard = newShard(true); } @Override public void tearDown() throws Exception { super.tearDown(); + if (executor != null) { + ThreadPool.terminate(executor, 10, TimeUnit.SECONDS); + } closeShards(indexShard); } @@ -121,12 +132,14 @@ public void testPostFilterDisablesCountOptimization() throws Exception { IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0)); + TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0, executor)); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers()); assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value); + // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery + // see: https://github.com/apache/lucene/pull/672 assertProfileData(context, "ConstantScoreQuery", query -> { assertThat(query.getTimeBreakdown().keySet(), not(empty())); assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); @@ -139,7 +152,7 @@ public void testPostFilterDisablesCountOptimization() throws Exception { assertThat(collector.getProfiledChildren(), empty()); }); - context.setSearcher(newContextSearcher(reader)); + context.setSearcher(newContextSearcher(reader, executor)); context.parsedPostFilter(new ParsedQuery(new MatchNoDocsQuery())); QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers()); assertEquals(0, context.queryResult().topDocs().topDocs.totalHits.value); @@ -157,6 +170,8 @@ public void testPostFilterDisablesCountOptimization() throws Exception { assertThat(query.getTimeBreakdown().get("create_weight"), greaterThan(0L)); assertThat(query.getTimeBreakdown().get("create_weight_count"), equalTo(1L)); }, (query) -> { + // IndexSearcher#rewrite optimizes by rewriting non-scoring queries to ConstantScoreQuery + // see: https://github.com/apache/lucene/pull/672 assertThat(query.getQueryName(), equalTo("ConstantScoreQuery")); assertThat(query.getTimeBreakdown().keySet(), not(empty())); assertThat(query.getTimeBreakdown().get("score"), equalTo(0L)); @@ -183,7 +198,7 @@ public void testTerminateAfterWithFilter() throws Exception { IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); @@ -233,7 +248,7 @@ public void testMinScoreDisablesCountOptimization() throws Exception { w.close(); IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0)); + TestSearchContext context = new TestSearchContext(null, indexShard, newEarlyTerminationContextSearcher(reader, 0, executor)); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); context.setSize(0); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); @@ -287,7 +302,7 @@ public void testInOrderScrollOptimization() throws Exception { w.close(); IndexReader reader = DirectoryReader.open(dir); ScrollContext scrollContext = new ScrollContext(); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader), scrollContext); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor), scrollContext); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); context.sort(new SortAndFormats(sort, new DocValueFormat[] { DocValueFormat.RAW })); scrollContext.lastEmittedDoc = null; @@ -300,6 +315,7 @@ public void testInOrderScrollOptimization() throws Exception { QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers()); assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); assertNull(context.queryResult().terminatedEarly()); + assertThat(context.terminateAfter(), equalTo(0)); assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); assertProfileData(context, "ConstantScoreQuery", query -> { assertThat(query.getTimeBreakdown().keySet(), not(empty())); @@ -313,7 +329,7 @@ public void testInOrderScrollOptimization() throws Exception { assertThat(collector.getProfiledChildren(), empty()); }); - context.setSearcher(newEarlyTerminationContextSearcher(reader, size)); + context.setSearcher(newEarlyTerminationContextSearcher(reader, size, executor)); QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers()); assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs)); @@ -356,7 +372,7 @@ public void testTerminateAfterEarlyTermination() throws Exception { } w.close(); final IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); @@ -573,7 +589,7 @@ public void testIndexSortingEarlyTermination() throws Exception { w.close(); final IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); context.setSize(1); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); @@ -634,7 +650,7 @@ public void testIndexSortingEarlyTermination() throws Exception { } { - context.setSearcher(newEarlyTerminationContextSearcher(reader, 1)); + context.setSearcher(newEarlyTerminationContextSearcher(reader, 1, executor)); context.trackTotalHitsUpTo(SearchContext.TRACK_TOTAL_HITS_DISABLED); QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers()); assertNull(context.queryResult().terminatedEarly()); @@ -703,7 +719,7 @@ public void testIndexSortScrollOptimization() throws Exception { searchSortAndFormats.add(new SortAndFormats(new Sort(indexSort.getSort()[0]), new DocValueFormat[] { DocValueFormat.RAW })); for (SortAndFormats searchSortAndFormat : searchSortAndFormats) { ScrollContext scrollContext = new ScrollContext(); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader), scrollContext); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor), scrollContext); context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery())); scrollContext.lastEmittedDoc = null; scrollContext.maxScore = Float.NaN; @@ -734,7 +750,7 @@ public void testIndexSortScrollOptimization() throws Exception { int sizeMinus1 = context.queryResult().topDocs().topDocs.scoreDocs.length - 1; FieldDoc lastDoc = (FieldDoc) context.queryResult().topDocs().topDocs.scoreDocs[sizeMinus1]; - context.setSearcher(newEarlyTerminationContextSearcher(reader, 10)); + context.setSearcher(newEarlyTerminationContextSearcher(reader, 10, executor)); QueryPhase.executeInternal(context.withCleanQueryResult().withProfilers()); assertNull(context.queryResult().terminatedEarly()); assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs)); @@ -764,7 +780,7 @@ public void testIndexSortScrollOptimization() throws Exception { @SuppressWarnings("unchecked") FieldComparator comparator = (FieldComparator) searchSortAndFormat.sort.getSort()[i].getComparator( i, - false + randomBoolean() ); int cmp = comparator.compareValues(firstDoc.fields[i], lastDoc.fields[i]); if (cmp == 0) { @@ -796,7 +812,7 @@ public void testDisableTopScoreCollection() throws Exception { w.close(); IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap())); Query q = new SpanNearQuery.Builder("title", true).addClause(new SpanTermQuery(new Term("title", "foo"))) .addClause(new SpanTermQuery(new Term("title", "bar"))) @@ -861,7 +877,7 @@ public void testMinScore() throws Exception { w.close(); IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.parsedQuery( new ParsedQuery( new BooleanQuery.Builder().add(new TermQuery(new Term("foo", "bar")), Occur.MUST) @@ -922,7 +938,7 @@ public void testMaxScore() throws Exception { w.close(); IndexReader reader = DirectoryReader.open(dir); - TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader, executor)); context.trackScores(true); context.parsedQuery( new ParsedQuery( @@ -1016,7 +1032,7 @@ public void testCollapseQuerySearchResults() throws Exception { new NumberFieldType("user", NumberType.INTEGER, true, false, true, false, null, Collections.emptyMap()) ); - TestSearchContext context = new TestSearchContext(queryShardContext, indexShard, newContextSearcher(reader)); + TestSearchContext context = new TestSearchContext(queryShardContext, indexShard, newContextSearcher(reader, executor)); context.collapse(new CollapseBuilder("user").build(context.getQueryShardContext())); context.trackScores(true); context.parsedQuery(new ParsedQuery(new TermQuery(new Term("foo", "bar")))); @@ -1129,25 +1145,26 @@ private final void assertProfileData( collector.accept(queryProfileShardResult.getCollectorResult()); } - private static ContextIndexSearcher newContextSearcher(IndexReader reader) throws IOException { + private static ContextIndexSearcher newContextSearcher(IndexReader reader, ExecutorService executor) throws IOException { return new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - null + executor ); } - private static ContextIndexSearcher newEarlyTerminationContextSearcher(IndexReader reader, int size) throws IOException { + private static ContextIndexSearcher newEarlyTerminationContextSearcher(IndexReader reader, int size, ExecutorService executor) + throws IOException { return new ContextIndexSearcher( reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true, - null + executor ) { @Override diff --git a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java index f218895754b7f..9fc81a500f6b0 100644 --- a/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java +++ b/server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java @@ -2071,8 +2071,7 @@ public void onFailure(final Exception e) { new QueryPhase(), new FetchPhase(Collections.emptyList()), responseCollectorService, - new NoneCircuitBreakerService(), - null + new NoneCircuitBreakerService() ); SearchPhaseController searchPhaseController = new SearchPhaseController( writableRegistry(), diff --git a/test/framework/src/main/java/org/opensearch/node/MockNode.java b/test/framework/src/main/java/org/opensearch/node/MockNode.java index 1bb034588d20e..6c70a6655a47f 100644 --- a/test/framework/src/main/java/org/opensearch/node/MockNode.java +++ b/test/framework/src/main/java/org/opensearch/node/MockNode.java @@ -72,7 +72,6 @@ import java.util.Collections; import java.util.Map; import java.util.Set; -import java.util.concurrent.Executor; import java.util.function.Function; /** @@ -153,8 +152,7 @@ protected SearchService newSearchService( QueryPhase queryPhase, FetchPhase fetchPhase, ResponseCollectorService responseCollectorService, - CircuitBreakerService circuitBreakerService, - Executor indexSearcherExecutor + CircuitBreakerService circuitBreakerService ) { if (getPluginsService().filterPlugins(MockSearchService.TestPlugin.class).isEmpty()) { return super.newSearchService( @@ -166,8 +164,7 @@ protected SearchService newSearchService( queryPhase, fetchPhase, responseCollectorService, - circuitBreakerService, - indexSearcherExecutor + circuitBreakerService ); } return new MockSearchService( diff --git a/test/framework/src/main/java/org/opensearch/search/MockSearchService.java b/test/framework/src/main/java/org/opensearch/search/MockSearchService.java index fd522cd298632..532b1fb844d52 100644 --- a/test/framework/src/main/java/org/opensearch/search/MockSearchService.java +++ b/test/framework/src/main/java/org/opensearch/search/MockSearchService.java @@ -96,18 +96,7 @@ public MockSearchService( FetchPhase fetchPhase, CircuitBreakerService circuitBreakerService ) { - super( - clusterService, - indicesService, - threadPool, - scriptService, - bigArrays, - queryPhase, - fetchPhase, - null, - circuitBreakerService, - null - ); + super(clusterService, indicesService, threadPool, scriptService, bigArrays, queryPhase, fetchPhase, null, circuitBreakerService); } @Override