org.apache.lucene.search.RandomApproximationQuery.java Source code

Introduction

Here is the source code for org.apache.lucene.search.RandomApproximationQuery.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.lucene.search;

import java.io.IOException;
import java.util.Random;

import com.carrotsearch.randomizedtesting.generators.RandomNumbers;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReaderContext;

/**
 * A {@link Query} that adds random approximations to its scorers.
 */
public class RandomApproximationQuery extends Query {

    private final Query query;
    private final Random random;

    public RandomApproximationQuery(Query query, Random random) {
        this.query = query;
        this.random = random;
    }

    @Override
    public Query rewrite(IndexReader reader) throws IOException {
        final Query rewritten = query.rewrite(reader);
        if (rewritten != query) {
            return new RandomApproximationQuery(rewritten, random);
        }
        return super.rewrite(reader);
    }

    @Override
    public void visit(QueryVisitor visitor) {
        query.visit(visitor);
    }

    @Override
    public boolean equals(Object other) {
        return sameClassAs(other) && query.equals(((RandomApproximationQuery) other).query);
    }

    @Override
    public int hashCode() {
        return 31 * classHash() + query.hashCode();
    }

    @Override
    public String toString(String field) {
        return query.toString(field);
    }

    @Override
    public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
        final Weight weight = query.createWeight(searcher, scoreMode, boost);
        return new RandomApproximationWeight(weight, new Random(random.nextLong()));
    }

    private static class RandomApproximationWeight extends FilterWeight {

        private final Random random;

        RandomApproximationWeight(Weight weight, Random random) {
            super(weight);
            this.random = random;
        }

        @Override
        public Scorer scorer(LeafReaderContext context) throws IOException {
            final Scorer scorer = in.scorer(context);
            if (scorer == null) {
                return null;
            }
            return new RandomApproximationScorer(scorer, new Random(random.nextLong()));
        }

    }

    private static class RandomApproximationScorer extends Scorer {

        private final Scorer scorer;
        private final RandomTwoPhaseView twoPhaseView;

        RandomApproximationScorer(Scorer scorer, Random random) {
            super(scorer.getWeight());
            this.scorer = scorer;
            this.twoPhaseView = new RandomTwoPhaseView(random, scorer.iterator());
        }

        @Override
        public TwoPhaseIterator twoPhaseIterator() {
            return twoPhaseView;
        }

        @Override
        public float score() throws IOException {
            return scorer.score();
        }

        @Override
        public int advanceShallow(int target) throws IOException {
            if (scorer.docID() > target && twoPhaseView.approximation.docID() != scorer.docID()) {
                // The random approximation can return doc ids that are not present in the underlying
                // scorer. These additional doc ids are always *before* the next matching doc so we
                // cannot use them to shallow advance the main scorer which is already ahead.
                target = scorer.docID();
            }
            return scorer.advanceShallow(target);
        }

        @Override
        public float getMaxScore(int upTo) throws IOException {
            return scorer.getMaxScore(upTo);
        }

        @Override
        public int docID() {
            return twoPhaseView.approximation().docID();
        }

        @Override
        public DocIdSetIterator iterator() {
            return TwoPhaseIterator.asDocIdSetIterator(twoPhaseView);
        }

    }

    private static class RandomTwoPhaseView extends TwoPhaseIterator {

        private final DocIdSetIterator disi;
        private int lastDoc = -1;
        private final float randomMatchCost;

        RandomTwoPhaseView(Random random, DocIdSetIterator disi) {
            super(new RandomApproximation(random, disi));
            this.disi = disi;
            this.randomMatchCost = random.nextFloat() * 200; // between 0 and 200
        }

        @Override
        public boolean matches() throws IOException {
            if (approximation.docID() == -1 || approximation.docID() == DocIdSetIterator.NO_MORE_DOCS) {
                throw new AssertionError("matches() should not be called on doc ID " + approximation.docID());
            }
            if (lastDoc == approximation.docID()) {
                throw new AssertionError("matches() has been called twice on doc ID " + approximation.docID());
            }
            lastDoc = approximation.docID();
            return approximation.docID() == disi.docID();
        }

        @Override
        public float matchCost() {
            return randomMatchCost;
        }
    }

    private static class RandomApproximation extends DocIdSetIterator {

        private final Random random;
        private final DocIdSetIterator disi;

        int doc = -1;

        public RandomApproximation(Random random, DocIdSetIterator disi) {
            this.random = random;
            this.disi = disi;
        }

        @Override
        public int docID() {
            return doc;
        }

        @Override
        public int nextDoc() throws IOException {
            return advance(doc + 1);
        }

        @Override
        public int advance(int target) throws IOException {
            if (disi.docID() < target) {
                disi.advance(target);
            }
            if (disi.docID() == NO_MORE_DOCS) {
                return doc = NO_MORE_DOCS;
            }
            return doc = RandomNumbers.randomIntBetween(random, target, disi.docID());
        }

        @Override
        public long cost() {
            return disi.cost();
        }
    }

}