Examples with SearchShardTask org.opensearch.action.search.SearchShardTask used on opensource projects

Example 6 with SearchShardTask

use of org.opensearch.action.search.SearchShardTask in project OpenSearch by opensearch-project.

the class QueryPhaseTests method testIndexSortScrollOptimization.

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(100, 200);
    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<SortAndFormats> 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), 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);
        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));
        QueryPhase.executeInternal(context);
        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<Object> comparator = (FieldComparator<Object>) searchSortAndFormat.sort.getSort()[i].getComparator(1, i);
            int cmp = comparator.compareValues(firstDoc.fields[i], lastDoc.fields[i]);
            if (cmp == 0) {
                continue;
            }
            assertThat(cmp, equalTo(1));
            break;
        }
    }
    reader.close();
    dir.close();
}

Example 7 with SearchShardTask

use of org.opensearch.action.search.SearchShardTask in project OpenSearch by opensearch-project.

the class QueryPhaseTests method testEnhanceSortOnNumeric.

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)));
        when(searchContext.mapperService()).thenReturn(mapperService);
        searchContext.sort(longSortAndFormats);
        searchContext.parsedQuery(query);
        searchContext.setTask(task);
        searchContext.setSize(10);
        QueryPhase.executeInternal(searchContext);
        assertTrue(searchContext.sort().sort.getSort()[0].getCanUsePoints());
        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)));
        when(searchContext.mapperService()).thenReturn(mapperService);
        searchContext.sort(longDateSortAndFormats);
        searchContext.parsedQuery(query);
        searchContext.setTask(task);
        searchContext.setSize(10);
        QueryPhase.executeInternal(searchContext);
        assertTrue(searchContext.sort().sort.getSort()[0].getCanUsePoints());
        assertSortResults(searchContext.queryResult().topDocs().topDocs, (long) numDocs, true);
    }
    // 3. Test a sort on date field
    {
        TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader)));
        when(searchContext.mapperService()).thenReturn(mapperService);
        searchContext.sort(dateSortAndFormats);
        searchContext.parsedQuery(query);
        searchContext.setTask(task);
        searchContext.setSize(10);
        QueryPhase.executeInternal(searchContext);
        assertTrue(searchContext.sort().sort.getSort()[0].getCanUsePoints());
        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)));
        when(searchContext.mapperService()).thenReturn(mapperService);
        searchContext.sort(dateLongSortAndFormats);
        searchContext.parsedQuery(query);
        searchContext.setTask(task);
        searchContext.setSize(10);
        QueryPhase.executeInternal(searchContext);
        assertTrue(searchContext.sort().sort.getSort()[0].getCanUsePoints());
        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)));
        when(searchContext.mapperService()).thenReturn(mapperService);
        searchContext.sort(longSortAndFormats);
        searchContext.parsedQuery(query);
        searchContext.setTask(task);
        searchContext.from(5);
        searchContext.setSize(0);
        QueryPhase.executeInternal(searchContext);
        assertTrue(searchContext.sort().sort.getSort()[0].getCanUsePoints());
        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)));
        when(searchContext.mapperService()).thenReturn(mapperService);
        searchContext.sort(longSortAndFormats);
        searchContext.parsedQuery(query);
        searchContext.setTask(task);
        searchContext.setSize(0);
        QueryPhase.executeInternal(searchContext);
    }
    // 7. Test that sort optimization works with search after
    {
        TestSearchContext searchContext = spy(new TestSearchContext(null, indexShard, newContextSearcher(reader)));
        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);
        assertTrue(searchContext.sort().sort.getSort()[0].getCanUsePoints());
        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);
    }
    reader.close();
    dir.close();
}

Example 8 with SearchShardTask

use of org.opensearch.action.search.SearchShardTask in project OpenSearch by opensearch-project.

the class QueryPhaseTests method testTerminateAfterWithFilter.

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));
    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);
        assertEquals(1, context.queryResult().topDocs().topDocs.totalHits.value);
        assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1));
    }
    reader.close();
    dir.close();
}

Example 9 with SearchShardTask

use of org.opensearch.action.search.SearchShardTask in project OpenSearch by opensearch-project.

the class QueryPhaseTests method testDisableTopScoreCollection.

public void testDisableTopScoreCollection() throws Exception {
    Directory dir = newDirectory();
    IndexWriterConfig iwc = newIndexWriterConfig(new StandardAnalyzer());
    RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc);
    Document doc = new Document();
    for (int i = 0; i < 10; 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));
    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);
    assertEquals(5, 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);
    assertEquals(5, 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();
}

Example 10 with SearchShardTask

use of org.opensearch.action.search.SearchShardTask in project OpenSearch by opensearch-project.

the class SearchSlowLogTests method testSlowLogsWithStats.

public void testSlowLogsWithStats() throws IOException {
    IndexService index = createIndex("foo");
    SearchContext searchContext = createSearchContext(index, "group1");
    SearchSourceBuilder source = SearchSourceBuilder.searchSource().query(QueryBuilders.matchAllQuery());
    searchContext.request().source(source);
    searchContext.setTask(new SearchShardTask(0, "n/a", "n/a", "test", null, Collections.singletonMap(Task.X_OPAQUE_ID, "my_id")));
    SearchSlowLog.SearchSlowLogMessage p = new SearchSlowLog.SearchSlowLogMessage(searchContext, 10);
    assertThat(p.getValueFor("stats"), equalTo("[\\\"group1\\\"]"));
    searchContext = createSearchContext(index, "group1", "group2");
    source = SearchSourceBuilder.searchSource().query(QueryBuilders.matchAllQuery());
    searchContext.request().source(source);
    searchContext.setTask(new SearchShardTask(0, "n/a", "n/a", "test", null, Collections.singletonMap(Task.X_OPAQUE_ID, "my_id")));
    p = new SearchSlowLog.SearchSlowLogMessage(searchContext, 10);
    assertThat(p.getValueFor("stats"), equalTo("[\\\"group1\\\", \\\"group2\\\"]"));
}