Examples with ParsedQuery org.opensearch.index.query.ParsedQuery used on opensource projects

Example 16 with ParsedQuery

use of org.opensearch.index.query.ParsedQuery in project OpenSearch by opensearch-project.

the class QueryPhaseTests method testTerminateAfterEarlyTermination.

public void testTerminateAfterEarlyTermination() throws Exception {
    Directory dir = newDirectory();
    IndexWriterConfig iwc = newIndexWriterConfig();
    RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc);
    final int numDocs = scaledRandomIntBetween(100, 200);
    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));
    context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap()));
    context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery()));
    context.terminateAfter(numDocs);
    {
        context.setSize(10);
        TotalHitCountCollector collector = new TotalHitCountCollector();
        context.queryCollectors().put(TotalHitCountCollector.class, collector);
        QueryPhase.executeInternal(context);
        assertFalse(context.queryResult().terminatedEarly());
        assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs));
        assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(10));
        assertThat(collector.getTotalHits(), equalTo(numDocs));
    }
    context.terminateAfter(1);
    {
        context.setSize(1);
        QueryPhase.executeInternal(context);
        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);
        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);
        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);
        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);
        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);
        TotalHitCountCollector collector = new TotalHitCountCollector();
        context.queryCollectors().put(TotalHitCountCollector.class, collector);
        QueryPhase.executeInternal(context);
        assertTrue(context.queryResult().terminatedEarly());
        assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L));
        assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1));
        assertThat(collector.getTotalHits(), equalTo(1));
        context.queryCollectors().clear();
    }
    {
        context.setSize(0);
        TotalHitCountCollector collector = new TotalHitCountCollector();
        context.queryCollectors().put(TotalHitCountCollector.class, collector);
        QueryPhase.executeInternal(context);
        assertTrue(context.queryResult().terminatedEarly());
        assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L));
        assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0));
        assertThat(collector.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);
        TotalHitCountCollector collector = new TotalHitCountCollector();
        context.queryCollectors().put(TotalHitCountCollector.class, collector);
        QueryPhase.executeInternal(context);
        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));
        assertThat(collector.getTotalHits(), equalTo(10));
    }
    context.terminateAfter(7);
    context.setSize(10);
    for (int trackTotalHits : new int[] { -1, 3, 75, 100 }) {
        context.trackTotalHitsUpTo(trackTotalHits);
        EarlyTerminatingCollector collector = new EarlyTerminatingCollector(new TotalHitCountCollector(), 1, false);
        context.queryCollectors().put(EarlyTerminatingCollector.class, collector);
        QueryPhase.executeInternal(context);
        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();
}

Example 17 with ParsedQuery

use of org.opensearch.index.query.ParsedQuery in project OpenSearch by opensearch-project.

the class QueryPhaseTests method countTestCase.

private void countTestCase(Query query, IndexReader reader, boolean shouldCollectSearch, boolean shouldCollectCount) throws Exception {
    ContextIndexSearcher searcher = shouldCollectSearch ? newContextSearcher(reader) : newEarlyTerminationContextSearcher(reader, 0);
    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);
    assertFalse(rescore);
    ContextIndexSearcher countSearcher = shouldCollectCount ? newContextSearcher(reader) : newEarlyTerminationContextSearcher(reader, 0);
    assertEquals(countSearcher.count(query), context.queryResult().topDocs().topDocs.totalHits.value);
}

Example 18 with ParsedQuery

use of org.opensearch.index.query.ParsedQuery in project OpenSearch by opensearch-project.

the class QueryPhaseTests method testPostFilterDisablesCountOptimization.

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));
    context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap()));
    context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery()));
    QueryPhase.executeInternal(context);
    assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value);
    context.setSearcher(newContextSearcher(reader));
    context.parsedPostFilter(new ParsedQuery(new MatchNoDocsQuery()));
    QueryPhase.executeInternal(context);
    assertEquals(0, context.queryResult().topDocs().topDocs.totalHits.value);
    reader.close();
    dir.close();
}

Example 19 with ParsedQuery

use of org.opensearch.index.query.ParsedQuery in project OpenSearch by opensearch-project.

the class QueryPhaseTests method testMinScoreDisablesCountOptimization.

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));
    context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery()));
    context.setSize(0);
    context.setTask(new SearchShardTask(123L, "", "", "", null, Collections.emptyMap()));
    QueryPhase.executeInternal(context);
    assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value);
    context.minimumScore(100);
    QueryPhase.executeInternal(context);
    assertEquals(0, context.queryResult().topDocs().topDocs.totalHits.value);
    reader.close();
    dir.close();
}

Example 20 with ParsedQuery

use of org.opensearch.index.query.ParsedQuery in project OpenSearch by opensearch-project.

the class DefaultSearchContextTests method testPreProcess.

public void testPreProcess() throws Exception {
    TimeValue timeout = new TimeValue(randomIntBetween(1, 100));
    ShardSearchRequest shardSearchRequest = mock(ShardSearchRequest.class);
    when(shardSearchRequest.searchType()).thenReturn(SearchType.DEFAULT);
    ShardId shardId = new ShardId("index", UUID.randomUUID().toString(), 1);
    when(shardSearchRequest.shardId()).thenReturn(shardId);
    ThreadPool threadPool = new TestThreadPool(this.getClass().getName());
    IndexShard indexShard = mock(IndexShard.class);
    QueryCachingPolicy queryCachingPolicy = mock(QueryCachingPolicy.class);
    when(indexShard.getQueryCachingPolicy()).thenReturn(queryCachingPolicy);
    when(indexShard.getThreadPool()).thenReturn(threadPool);
    int maxResultWindow = randomIntBetween(50, 100);
    int maxRescoreWindow = randomIntBetween(50, 100);
    int maxSlicesPerScroll = randomIntBetween(50, 100);
    Settings settings = Settings.builder().put("index.max_result_window", maxResultWindow).put("index.max_slices_per_scroll", maxSlicesPerScroll).put("index.max_rescore_window", maxRescoreWindow).put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1).put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 2).build();
    IndexService indexService = mock(IndexService.class);
    IndexCache indexCache = mock(IndexCache.class);
    QueryCache queryCache = mock(QueryCache.class);
    when(indexCache.query()).thenReturn(queryCache);
    when(indexService.cache()).thenReturn(indexCache);
    QueryShardContext queryShardContext = mock(QueryShardContext.class);
    when(indexService.newQueryShardContext(eq(shardId.id()), any(), any(), nullable(String.class))).thenReturn(queryShardContext);
    MapperService mapperService = mock(MapperService.class);
    when(mapperService.hasNested()).thenReturn(randomBoolean());
    when(indexService.mapperService()).thenReturn(mapperService);
    IndexMetadata indexMetadata = IndexMetadata.builder("index").settings(settings).build();
    IndexSettings indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY);
    when(indexService.getIndexSettings()).thenReturn(indexSettings);
    when(mapperService.getIndexSettings()).thenReturn(indexSettings);
    BigArrays bigArrays = new MockBigArrays(new MockPageCacheRecycler(Settings.EMPTY), new NoneCircuitBreakerService());
    try (Directory dir = newDirectory();
        RandomIndexWriter w = new RandomIndexWriter(random(), dir)) {
        final Supplier<Engine.SearcherSupplier> searcherSupplier = () -> new Engine.SearcherSupplier(Function.identity()) {

            @Override
            protected void doClose() {
            }

            @Override
            protected Engine.Searcher acquireSearcherInternal(String source) {
                try {
                    IndexReader reader = w.getReader();
                    return new Engine.Searcher("test", reader, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), reader);
                } catch (IOException exc) {
                    throw new AssertionError(exc);
                }
            }
        };
        SearchShardTarget target = new SearchShardTarget("node", shardId, null, OriginalIndices.NONE);
        ReaderContext readerWithoutScroll = new ReaderContext(newContextId(), indexService, indexShard, searcherSupplier.get(), randomNonNegativeLong(), false);
        DefaultSearchContext contextWithoutScroll = new DefaultSearchContext(readerWithoutScroll, shardSearchRequest, target, null, bigArrays, null, timeout, null, false, Version.CURRENT);
        contextWithoutScroll.from(300);
        contextWithoutScroll.close();
        // resultWindow greater than maxResultWindow and scrollContext is null
        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> contextWithoutScroll.preProcess(false));
        assertThat(exception.getMessage(), equalTo("Result window is too large, from + size must be less than or equal to:" + " [" + maxResultWindow + "] but was [310]. See the scroll api for a more efficient way to request large data sets. " + "This limit can be set by changing the [" + IndexSettings.MAX_RESULT_WINDOW_SETTING.getKey() + "] index level setting."));
        // resultWindow greater than maxResultWindow and scrollContext isn't null
        when(shardSearchRequest.scroll()).thenReturn(new Scroll(TimeValue.timeValueMillis(randomInt(1000))));
        ReaderContext readerContext = new LegacyReaderContext(newContextId(), indexService, indexShard, searcherSupplier.get(), shardSearchRequest, randomNonNegativeLong());
        DefaultSearchContext context1 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null, bigArrays, null, timeout, null, false, Version.CURRENT);
        context1.from(300);
        exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
        assertThat(exception.getMessage(), equalTo("Batch size is too large, size must be less than or equal to: [" + maxResultWindow + "] but was [310]. Scroll batch sizes cost as much memory as result windows so they are " + "controlled by the [" + IndexSettings.MAX_RESULT_WINDOW_SETTING.getKey() + "] index level setting."));
        // resultWindow not greater than maxResultWindow and both rescore and sort are not null
        context1.from(0);
        DocValueFormat docValueFormat = mock(DocValueFormat.class);
        SortAndFormats sortAndFormats = new SortAndFormats(new Sort(), new DocValueFormat[] { docValueFormat });
        context1.sort(sortAndFormats);
        RescoreContext rescoreContext = mock(RescoreContext.class);
        when(rescoreContext.getWindowSize()).thenReturn(500);
        context1.addRescore(rescoreContext);
        exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
        assertThat(exception.getMessage(), equalTo("Cannot use [sort] option in conjunction with [rescore]."));
        // rescore is null but sort is not null and rescoreContext.getWindowSize() exceeds maxResultWindow
        context1.sort(null);
        exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
        assertThat(exception.getMessage(), equalTo("Rescore window [" + rescoreContext.getWindowSize() + "] is too large. " + "It must be less than [" + maxRescoreWindow + "]. This prevents allocating massive heaps for storing the results " + "to be rescored. This limit can be set by changing the [" + IndexSettings.MAX_RESCORE_WINDOW_SETTING.getKey() + "] index level setting."));
        readerContext.close();
        readerContext = new ReaderContext(newContextId(), indexService, indexShard, searcherSupplier.get(), randomNonNegativeLong(), false);
        // rescore is null but sliceBuilder is not null
        DefaultSearchContext context2 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null, bigArrays, null, timeout, null, false, Version.CURRENT);
        SliceBuilder sliceBuilder = mock(SliceBuilder.class);
        int numSlices = maxSlicesPerScroll + randomIntBetween(1, 100);
        when(sliceBuilder.getMax()).thenReturn(numSlices);
        context2.sliceBuilder(sliceBuilder);
        exception = expectThrows(IllegalArgumentException.class, () -> context2.preProcess(false));
        assertThat(exception.getMessage(), equalTo("The number of slices [" + numSlices + "] is too large. It must " + "be less than [" + maxSlicesPerScroll + "]. This limit can be set by changing the [" + IndexSettings.MAX_SLICES_PER_SCROLL.getKey() + "] index level setting."));
        // No exceptions should be thrown
        when(shardSearchRequest.getAliasFilter()).thenReturn(AliasFilter.EMPTY);
        when(shardSearchRequest.indexBoost()).thenReturn(AbstractQueryBuilder.DEFAULT_BOOST);
        DefaultSearchContext context3 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null, bigArrays, null, timeout, null, false, Version.CURRENT);
        ParsedQuery parsedQuery = ParsedQuery.parsedMatchAllQuery();
        context3.sliceBuilder(null).parsedQuery(parsedQuery).preProcess(false);
        assertEquals(context3.query(), context3.buildFilteredQuery(parsedQuery.query()));
        when(queryShardContext.getIndexSettings()).thenReturn(indexSettings);
        when(queryShardContext.fieldMapper(anyString())).thenReturn(mock(MappedFieldType.class));
        when(shardSearchRequest.indexRoutings()).thenReturn(new String[0]);
        readerContext.close();
        readerContext = new ReaderContext(newContextId(), indexService, indexShard, searcherSupplier.get(), randomNonNegativeLong(), false);
        DefaultSearchContext context4 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null, bigArrays, null, timeout, null, false, Version.CURRENT);
        context4.sliceBuilder(new SliceBuilder(1, 2)).parsedQuery(parsedQuery).preProcess(false);
        Query query1 = context4.query();
        context4.sliceBuilder(new SliceBuilder(0, 2)).parsedQuery(parsedQuery).preProcess(false);
        Query query2 = context4.query();
        assertTrue(query1 instanceof MatchNoDocsQuery || query2 instanceof MatchNoDocsQuery);
        readerContext.close();
        threadPool.shutdown();
    }
}