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

Example 6 with ParsedQuery

use of org.opensearch.index.query.ParsedQuery 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 7 with ParsedQuery

use of org.opensearch.index.query.ParsedQuery 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 8 with ParsedQuery

use of org.opensearch.index.query.ParsedQuery 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 9 with ParsedQuery

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

the class DlsQueryParser method parse.

static ParsedQuery parse(final Set<String> unparsedDlsQueries, ParsedQuery originalQuery, final QueryShardContext queryShardContext, final NamedXContentRegistry namedXContentRegistry) throws IOException {
    if (unparsedDlsQueries == null || unparsedDlsQueries.isEmpty()) {
        return null;
    }
    final boolean hasNestedMapping = queryShardContext.getMapperService().hasNested();
    BooleanQuery.Builder dlsQueryBuilder = new BooleanQuery.Builder();
    dlsQueryBuilder.setMinimumNumberShouldMatch(1);
    for (final String unparsedDlsQuery : unparsedDlsQueries) {
        try {
            final QueryBuilder qb = queries.get(unparsedDlsQuery, new Callable<QueryBuilder>() {

                @Override
                public QueryBuilder call() throws Exception {
                    final XContentParser parser = JsonXContent.jsonXContent.createParser(namedXContentRegistry, THROW_UNSUPPORTED_OPERATION, unparsedDlsQuery);
                    final QueryBuilder qb = AbstractQueryBuilder.parseInnerQueryBuilder(parser);
                    return qb;
                }
            });
            final ParsedQuery parsedQuery = queryShardContext.toQuery(qb);
            // no need for scoring here, so its possible to wrap this in a
            // ConstantScoreQuery
            final Query dlsQuery = new ConstantScoreQuery(parsedQuery.query());
            dlsQueryBuilder.add(dlsQuery, Occur.SHOULD);
            if (hasNestedMapping) {
                handleNested(queryShardContext, dlsQueryBuilder, dlsQuery);
            }
        } catch (ExecutionException e) {
            throw new IOException(e);
        }
    }
    dlsQueryBuilder.add(originalQuery.query(), Occur.MUST);
    return new ParsedQuery(dlsQueryBuilder.build());
}

Example 10 with ParsedQuery

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

the class DlsQueryParser method parse.

static Query parse(final Set<String> unparsedDlsQueries, final QueryShardContext queryShardContext, final NamedXContentRegistry namedXContentRegistry) throws IOException {
    if (unparsedDlsQueries == null || unparsedDlsQueries.isEmpty()) {
        return null;
    }
    final boolean hasNestedMapping = queryShardContext.getMapperService().hasNested();
    BooleanQuery.Builder dlsQueryBuilder = new BooleanQuery.Builder();
    dlsQueryBuilder.setMinimumNumberShouldMatch(1);
    for (final String unparsedDlsQuery : unparsedDlsQueries) {
        try {
            final QueryBuilder qb = queries.get(unparsedDlsQuery, new Callable<QueryBuilder>() {

                @Override
                public QueryBuilder call() throws Exception {
                    final XContentParser parser = JsonXContent.jsonXContent.createParser(namedXContentRegistry, THROW_UNSUPPORTED_OPERATION, unparsedDlsQuery);
                    final QueryBuilder qb = AbstractQueryBuilder.parseInnerQueryBuilder(parser);
                    return qb;
                }
            });
            final ParsedQuery parsedQuery = queryShardContext.toQuery(qb);
            final Query dlsQuery = parsedQuery.query();
            dlsQueryBuilder.add(dlsQuery, Occur.SHOULD);
            if (hasNestedMapping) {
                handleNested(queryShardContext, dlsQueryBuilder, dlsQuery);
            }
        } catch (ExecutionException e) {
            throw new IOException(e);
        }
    }
    // ConstantScoreQuery
    return new ConstantScoreQuery(dlsQueryBuilder.build());
}