All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.apache.lucene.queries.SearchAfterSortedDocQuery Maven / Gradle / Ivy

Go to download

Elasticsearch subproject :distribution:archives:integ-test-zip

There is a newer version: 8.13.4
Show newest version
/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0 and the Server Side Public License, v 1; you may not use this file except
 * in compliance with, at your election, the Elastic License 2.0 or the Server
 * Side Public License, v 1.
 */

package org.apache.lucene.queries;

import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.ConstantScoreScorer;
import org.apache.lucene.search.ConstantScoreWeight;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.FieldComparator;
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.LeafFieldComparator;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.SortField;
import org.apache.lucene.search.Weight;
import org.elasticsearch.common.lucene.Lucene;

import java.io.IOException;
import java.util.Arrays;
import java.util.Objects;

/**
 * A {@link Query} that only matches documents that are greater than the provided {@link FieldDoc}.
 * This works only if the index is sorted according to the given search {@link Sort}.
 */
public class SearchAfterSortedDocQuery extends Query {
    private final Sort sort;
    private final FieldDoc after;
    private final FieldComparator[] fieldComparators;
    private final int[] reverseMuls;

    public SearchAfterSortedDocQuery(Sort sort, FieldDoc after) {
        if (sort.getSort().length != after.fields.length) {
            throw new IllegalArgumentException(
                "after doc  has " + after.fields.length + " value(s) but sort has " + sort.getSort().length + "."
            );
        }
        this.sort = Objects.requireNonNull(sort);
        this.after = after;
        int numFields = sort.getSort().length;
        this.fieldComparators = new FieldComparator[numFields];
        this.reverseMuls = new int[numFields];
        for (int i = 0; i < numFields; i++) {
            SortField sortField = sort.getSort()[i];
            FieldComparator fieldComparator = sortField.getComparator(1, i);
            @SuppressWarnings("unchecked")
            FieldComparator comparator = (FieldComparator) fieldComparator;
            comparator.setTopValue(after.fields[i]);
            fieldComparators[i] = fieldComparator;
            reverseMuls[i] = sortField.getReverse() ? -1 : 1;
        }
    }

    @Override
    public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
        return new ConstantScoreWeight(this, 1.0f) {
            @Override
            public Scorer scorer(LeafReaderContext context) throws IOException {
                Sort segmentSort = context.reader().getMetaData().getSort();
                if (segmentSort == null || Lucene.canEarlyTerminate(sort, segmentSort) == false) {
                    throw new IOException("search sort :[" + sort + "] does not match the index sort:[" + segmentSort + "]");
                }
                final int afterDoc = after.doc - context.docBase;
                TopComparator comparator = getTopComparator(fieldComparators, reverseMuls, context, afterDoc);
                final int maxDoc = context.reader().maxDoc();
                final int firstDoc = searchAfterDoc(comparator, 0, context.reader().maxDoc());
                if (firstDoc >= maxDoc) {
                    return null;
                }
                final DocIdSetIterator disi = new MinDocQuery.MinDocIterator(firstDoc, maxDoc);
                return new ConstantScoreScorer(this, score(), scoreMode, disi);
            }

            @Override
            public boolean isCacheable(LeafReaderContext ctx) {
                // If the sort order includes _doc, then the matches in a segment
                // may depend on other segments, which makes this query a bad
                // candidate for caching
                return false;
            }
        };
    }

    @Override
    public String toString(String field) {
        return "SearchAfterSortedDocQuery(sort=" + sort + ", afterDoc=" + after.toString() + ")";
    }

    @Override
    public boolean equals(Object other) {
        return sameClassAs(other) && equalsTo(getClass().cast(other));
    }

    private boolean equalsTo(SearchAfterSortedDocQuery other) {
        return sort.equals(other.sort)
            && after.doc == other.after.doc
            && Double.compare(after.score, other.after.score) == 0
            && Arrays.equals(after.fields, other.after.fields);
    }

    @Override
    public int hashCode() {
        return Objects.hash(classHash(), sort, after.doc, after.score, Arrays.hashCode(after.fields));
    }

    interface TopComparator {
        boolean lessThanTop(int doc) throws IOException;
    }

    static TopComparator getTopComparator(
        FieldComparator[] fieldComparators,
        int[] reverseMuls,
        LeafReaderContext leafReaderContext,
        int topDoc
    ) {
        return doc -> {
            // DVs use forward iterators so we recreate the iterator for each sort field
            // every time we need to compare a document with the after doc.
            // We could reuse the iterators when the comparison goes forward but
            // this should only be called a few time per segment (binary search).
            for (int i = 0; i < fieldComparators.length; i++) {
                LeafFieldComparator comparator = fieldComparators[i].getLeafComparator(leafReaderContext);
                int value = reverseMuls[i] * comparator.compareTop(doc);
                if (value != 0) {
                    return value < 0;
                }
            }

            if (doc <= topDoc) {
                return false;
            }
            return true;
        };
    }

    /**
     * Returns the first doc id greater than the provided after doc.
     */
    static int searchAfterDoc(TopComparator comparator, int from, int to) throws IOException {
        int low = from;
        int high = to - 1;

        while (low <= high) {
            int mid = (low + high) >>> 1;
            if (comparator.lessThanTop(mid)) {
                high = mid - 1;
            } else {
                low = mid + 1;
            }
        }
        return low;
    }

}
 















© 2015 - 2024 Weber Informatics LLC | Privacy Policy