Examples with TestSearchContext org.opensearch.test.TestSearchContext used on opensource projects

Example 1 with TestSearchContext

use of org.opensearch.test.TestSearchContext in project OpenSearch by opensearch-project.

the class AdjacencyMatrixAggregationBuilderTests method testFilterSizeLimitation.

public void testFilterSizeLimitation() throws Exception {
    // filter size grater than max size should thrown a exception
    QueryShardContext queryShardContext = mock(QueryShardContext.class);
    IndexShard indexShard = mock(IndexShard.class);
    Settings settings = Settings.builder().put("index.max_adjacency_matrix_filters", 2).put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1).put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 2).build();
    IndexMetadata indexMetadata = IndexMetadata.builder("index").settings(settings).build();
    IndexSettings indexSettings = new IndexSettings(indexMetadata, Settings.EMPTY);
    when(indexShard.indexSettings()).thenReturn(indexSettings);
    when(queryShardContext.getIndexSettings()).thenReturn(indexSettings);
    SearchContext context = new TestSearchContext(queryShardContext, indexShard);
    Map<String, QueryBuilder> filters = new HashMap<>(3);
    for (int i = 0; i < 3; i++) {
        QueryBuilder queryBuilder = mock(QueryBuilder.class);
        // return builder itself to skip rewrite
        when(queryBuilder.rewrite(queryShardContext)).thenReturn(queryBuilder);
        filters.put("filter" + i, queryBuilder);
    }
    AdjacencyMatrixAggregationBuilder builder = new AdjacencyMatrixAggregationBuilder("dummy", filters);
    IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> builder.doBuild(context.getQueryShardContext(), null, new AggregatorFactories.Builder()));
    assertThat(ex.getMessage(), equalTo("Number of filters is too large, must be less than or equal to: [2] but was [3]." + "This limit can be set by changing the [" + IndexSettings.MAX_ADJACENCY_MATRIX_FILTERS_SETTING.getKey() + "] index level setting."));
    // filter size not grater than max size should return an instance of AdjacencyMatrixAggregatorFactory
    Map<String, QueryBuilder> emptyFilters = Collections.emptyMap();
    AdjacencyMatrixAggregationBuilder aggregationBuilder = new AdjacencyMatrixAggregationBuilder("dummy", emptyFilters);
    AggregatorFactory factory = aggregationBuilder.doBuild(context.getQueryShardContext(), null, new AggregatorFactories.Builder());
    assertThat(factory instanceof AdjacencyMatrixAggregatorFactory, is(true));
    assertThat(factory.name(), equalTo("dummy"));
    assertWarnings("[index.max_adjacency_matrix_filters] setting was deprecated in OpenSearch and will be " + "removed in a future release! See the breaking changes documentation for the next major version.");
}

Example 2 with TestSearchContext

use of org.opensearch.test.TestSearchContext in project OpenSearch by opensearch-project.

the class QueryPhaseTests method testMinScore.

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));
    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);
    assertEquals(10, context.queryResult().topDocs().topDocs.totalHits.value);
    reader.close();
    dir.close();
}

Example 3 with TestSearchContext

use of org.opensearch.test.TestSearchContext in project OpenSearch by opensearch-project.

the class QueryPhaseTests method testInOrderScrollOptimization.

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(100, 200);
    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), 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()));
    int size = randomIntBetween(2, 5);
    context.setSize(size);
    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));
    context.setSearcher(newEarlyTerminationContextSearcher(reader, size));
    QueryPhase.executeInternal(context);
    assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) numDocs));
    assertThat(context.terminateAfter(), equalTo(size));
    assertThat(context.queryResult().getTotalHits().value, equalTo((long) numDocs));
    assertThat(context.queryResult().topDocs().topDocs.scoreDocs[0].doc, greaterThanOrEqualTo(size));
    reader.close();
    dir.close();
}

Example 4 with TestSearchContext

use of org.opensearch.test.TestSearchContext in project OpenSearch by opensearch-project.

the class QueryPhaseTests method testIndexSortingEarlyTermination.

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(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.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);
    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);
        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 TotalHitCountCollector totalHitCountCollector = new TotalHitCountCollector();
        context.queryCollectors().put(TotalHitCountCollector.class, totalHitCountCollector);
        QueryPhase.executeInternal(context);
        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)));
        assertThat(totalHitCountCollector.getTotalHits(), equalTo(numDocs));
        context.queryCollectors().clear();
    }
    {
        context.setSearcher(newEarlyTerminationContextSearcher(reader, 1));
        context.trackTotalHitsUpTo(SearchContext.TRACK_TOTAL_HITS_DISABLED);
        QueryPhase.executeInternal(context);
        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);
        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();
}

Example 5 with TestSearchContext

use of org.opensearch.test.TestSearchContext 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();
}