Examples with SortAndFormats org.opensearch.search.sort.SortAndFormats used on opensource projects

Example 1 with SortAndFormats

use of org.opensearch.search.sort.SortAndFormats in project OpenSearch by opensearch-project.

the class CompositeAggregator method processLeafFromQuery.

private void processLeafFromQuery(LeafReaderContext ctx, Sort indexSortPrefix) throws IOException {
    DocValueFormat[] formats = new DocValueFormat[indexSortPrefix.getSort().length];
    for (int i = 0; i < formats.length; i++) {
        formats[i] = sources[i].format;
    }
    FieldDoc fieldDoc = SearchAfterBuilder.buildFieldDoc(new SortAndFormats(indexSortPrefix, formats), Arrays.copyOfRange(rawAfterKey.values(), 0, formats.length));
    if (indexSortPrefix.getSort().length < sources.length) {
        // include all docs that belong to the partial bucket
        fieldDoc.doc = -1;
    }
    BooleanQuery newQuery = new BooleanQuery.Builder().add(context.query(), BooleanClause.Occur.MUST).add(new SearchAfterSortedDocQuery(applySortFieldRounding(indexSortPrefix), fieldDoc), BooleanClause.Occur.FILTER).build();
    Weight weight = context.searcher().createWeight(context.searcher().rewrite(newQuery), ScoreMode.COMPLETE_NO_SCORES, 1f);
    Scorer scorer = weight.scorer(ctx);
    if (scorer != null) {
        DocIdSetIterator docIt = scorer.iterator();
        final LeafBucketCollector inner = queue.getLeafCollector(ctx, getFirstPassCollector(docIdSetBuilder, indexSortPrefix.getSort().length));
        inner.setScorer(scorer);
        final Bits liveDocs = ctx.reader().getLiveDocs();
        while (docIt.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
            if (liveDocs == null || liveDocs.get(docIt.docID())) {
                inner.collect(docIt.docID());
            }
        }
    }
}

Example 2 with SortAndFormats

use of org.opensearch.search.sort.SortAndFormats in project OpenSearch by opensearch-project.

the class InnerHitContextBuilder method setupInnerHitsContext.

protected void setupInnerHitsContext(QueryShardContext queryShardContext, InnerHitsContext.InnerHitSubContext innerHitsContext) throws IOException {
    innerHitsContext.from(innerHitBuilder.getFrom());
    innerHitsContext.size(innerHitBuilder.getSize());
    innerHitsContext.explain(innerHitBuilder.isExplain());
    innerHitsContext.version(innerHitBuilder.isVersion());
    innerHitsContext.seqNoAndPrimaryTerm(innerHitBuilder.isSeqNoAndPrimaryTerm());
    innerHitsContext.trackScores(innerHitBuilder.isTrackScores());
    if (innerHitBuilder.getStoredFieldsContext() != null) {
        innerHitsContext.storedFieldsContext(innerHitBuilder.getStoredFieldsContext());
    }
    if (innerHitBuilder.getDocValueFields() != null) {
        FetchDocValuesContext docValuesContext = FetchDocValuesContext.create(queryShardContext.getMapperService(), innerHitBuilder.getDocValueFields());
        innerHitsContext.docValuesContext(docValuesContext);
    }
    if (innerHitBuilder.getFetchFields() != null) {
        String indexName = queryShardContext.index().getName();
        FetchFieldsContext fieldsContext = new FetchFieldsContext(innerHitBuilder.getFetchFields());
        innerHitsContext.fetchFieldsContext(fieldsContext);
    }
    if (innerHitBuilder.getScriptFields() != null) {
        for (SearchSourceBuilder.ScriptField field : innerHitBuilder.getScriptFields()) {
            QueryShardContext innerContext = innerHitsContext.getQueryShardContext();
            FieldScript.Factory factory = innerContext.compile(field.script(), FieldScript.CONTEXT);
            FieldScript.LeafFactory fieldScript = factory.newFactory(field.script().getParams(), innerContext.lookup());
            innerHitsContext.scriptFields().add(new org.opensearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(field.fieldName(), fieldScript, field.ignoreFailure()));
        }
    }
    if (innerHitBuilder.getFetchSourceContext() != null) {
        innerHitsContext.fetchSourceContext(innerHitBuilder.getFetchSourceContext());
    }
    if (innerHitBuilder.getSorts() != null) {
        Optional<SortAndFormats> optionalSort = SortBuilder.buildSort(innerHitBuilder.getSorts(), queryShardContext);
        if (optionalSort.isPresent()) {
            innerHitsContext.sort(optionalSort.get());
        }
    }
    if (innerHitBuilder.getHighlightBuilder() != null) {
        innerHitsContext.highlight(innerHitBuilder.getHighlightBuilder().build(queryShardContext));
    }
    ParsedQuery parsedQuery = new ParsedQuery(query.toQuery(queryShardContext), queryShardContext.copyNamedQueries());
    innerHitsContext.parsedQuery(parsedQuery);
    Map<String, InnerHitsContext.InnerHitSubContext> baseChildren = buildChildInnerHits(innerHitsContext.parentSearchContext(), children);
    innerHitsContext.setChildInnerHits(baseChildren);
}

Example 3 with SortAndFormats

use of org.opensearch.search.sort.SortAndFormats in project OpenSearch by opensearch-project.

the class TopHitsAggregationBuilder method doBuild.

@Override
protected TopHitsAggregatorFactory doBuild(QueryShardContext queryShardContext, AggregatorFactory parent, Builder subfactoriesBuilder) throws IOException {
    long innerResultWindow = from() + size();
    int maxInnerResultWindow = queryShardContext.getMapperService().getIndexSettings().getMaxInnerResultWindow();
    if (innerResultWindow > maxInnerResultWindow) {
        throw new IllegalArgumentException("Top hits result window is too large, the top hits aggregator [" + name + "]'s from + size must be less " + "than or equal to: [" + maxInnerResultWindow + "] but was [" + innerResultWindow + "]. This limit can be set by changing the [" + IndexSettings.MAX_INNER_RESULT_WINDOW_SETTING.getKey() + "] index level setting.");
    }
    List<ScriptFieldsContext.ScriptField> scriptFields = new ArrayList<>();
    if (this.scriptFields != null) {
        for (ScriptField field : this.scriptFields) {
            FieldScript.Factory factory = queryShardContext.compile(field.script(), FieldScript.CONTEXT);
            FieldScript.LeafFactory searchScript = factory.newFactory(field.script().getParams(), queryShardContext.lookup());
            scriptFields.add(new org.opensearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(field.fieldName(), searchScript, field.ignoreFailure()));
        }
    }
    final Optional<SortAndFormats> optionalSort;
    if (sorts == null) {
        optionalSort = Optional.empty();
    } else {
        optionalSort = SortBuilder.buildSort(sorts, queryShardContext);
    }
    return new TopHitsAggregatorFactory(name, from, size, explain, version, seqNoAndPrimaryTerm, trackScores, optionalSort, highlightBuilder, storedFieldsContext, docValueFields, fetchFields, scriptFields, fetchSourceContext, queryShardContext, parent, subfactoriesBuilder, metadata);
}

Example 4 with SortAndFormats

use of org.opensearch.search.sort.SortAndFormats 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 SortAndFormats

use of org.opensearch.search.sort.SortAndFormats 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();
}