Examples with Weight org.apache.lucene.search.Weight used on opensource projects

Example 1 with Weight

use of org.apache.lucene.search.Weight in project elasticsearch by elastic.

the class CompletionSuggester method suggest.

private static void suggest(IndexSearcher searcher, CompletionQuery query, TopSuggestDocsCollector collector) throws IOException {
    query = (CompletionQuery) query.rewrite(searcher.getIndexReader());
    Weight weight = query.createWeight(searcher, collector.needsScores());
    for (LeafReaderContext context : searcher.getIndexReader().leaves()) {
        BulkScorer scorer = weight.bulkScorer(context);
        if (scorer != null) {
            try {
                scorer.score(collector.getLeafCollector(context), context.reader().getLiveDocs());
            } catch (CollectionTerminatedException e) {
            // collection was terminated prematurely
            // continue with the following leaf
            }
        }
    }
}

Example 2 with Weight

use of org.apache.lucene.search.Weight in project elasticsearch by elastic.

the class CollapsingTopDocsCollectorTests method assertSearchCollapse.

private <T extends Comparable> void assertSearchCollapse(CollapsingDocValuesProducer<T> dvProducers, boolean numeric, boolean multivalued) throws IOException {
    final int numDocs = randomIntBetween(1000, 2000);
    int maxGroup = randomIntBetween(2, 500);
    final Directory dir = newDirectory();
    final RandomIndexWriter w = new RandomIndexWriter(random(), dir);
    Set<T> values = new HashSet<>();
    int totalHits = 0;
    for (int i = 0; i < numDocs; i++) {
        final T value = dvProducers.randomGroup(maxGroup);
        values.add(value);
        Document doc = new Document();
        dvProducers.add(doc, value, multivalued);
        doc.add(new NumericDocValuesField("sort1", randomIntBetween(0, 10)));
        doc.add(new NumericDocValuesField("sort2", randomLong()));
        w.addDocument(doc);
        totalHits++;
    }
    List<T> valueList = new ArrayList<>(values);
    Collections.sort(valueList);
    final IndexReader reader = w.getReader();
    final IndexSearcher searcher = newSearcher(reader);
    final SortField collapseField = dvProducers.sortField(multivalued);
    final SortField sort1 = new SortField("sort1", SortField.Type.INT);
    final SortField sort2 = new SortField("sort2", SortField.Type.LONG);
    Sort sort = new Sort(sort1, sort2, collapseField);
    int expectedNumGroups = values.size();
    final CollapsingTopDocsCollector collapsingCollector;
    if (numeric) {
        collapsingCollector = CollapsingTopDocsCollector.createNumeric(collapseField.getField(), sort, expectedNumGroups, false);
    } else {
        collapsingCollector = CollapsingTopDocsCollector.createKeyword(collapseField.getField(), sort, expectedNumGroups, false);
    }
    TopFieldCollector topFieldCollector = TopFieldCollector.create(sort, totalHits, true, false, false);
    searcher.search(new MatchAllDocsQuery(), collapsingCollector);
    searcher.search(new MatchAllDocsQuery(), topFieldCollector);
    CollapseTopFieldDocs collapseTopFieldDocs = collapsingCollector.getTopDocs();
    TopFieldDocs topDocs = topFieldCollector.topDocs();
    assertEquals(collapseField.getField(), collapseTopFieldDocs.field);
    assertEquals(expectedNumGroups, collapseTopFieldDocs.scoreDocs.length);
    assertEquals(totalHits, collapseTopFieldDocs.totalHits);
    assertEquals(totalHits, topDocs.scoreDocs.length);
    assertEquals(totalHits, topDocs.totalHits);
    Set<Object> seen = new HashSet<>();
    // collapse field is the last sort
    int collapseIndex = sort.getSort().length - 1;
    int topDocsIndex = 0;
    for (int i = 0; i < expectedNumGroups; i++) {
        FieldDoc fieldDoc = null;
        for (; topDocsIndex < totalHits; topDocsIndex++) {
            fieldDoc = (FieldDoc) topDocs.scoreDocs[topDocsIndex];
            if (seen.contains(fieldDoc.fields[collapseIndex]) == false) {
                break;
            }
        }
        FieldDoc collapseFieldDoc = (FieldDoc) collapseTopFieldDocs.scoreDocs[i];
        assertNotNull(fieldDoc);
        assertEquals(collapseFieldDoc.doc, fieldDoc.doc);
        assertArrayEquals(collapseFieldDoc.fields, fieldDoc.fields);
        seen.add(fieldDoc.fields[fieldDoc.fields.length - 1]);
    }
    for (; topDocsIndex < totalHits; topDocsIndex++) {
        FieldDoc fieldDoc = (FieldDoc) topDocs.scoreDocs[topDocsIndex];
        assertTrue(seen.contains(fieldDoc.fields[collapseIndex]));
    }
    // check merge
    final IndexReaderContext ctx = searcher.getTopReaderContext();
    final SegmentSearcher[] subSearchers;
    final int[] docStarts;
    if (ctx instanceof LeafReaderContext) {
        subSearchers = new SegmentSearcher[1];
        docStarts = new int[1];
        subSearchers[0] = new SegmentSearcher((LeafReaderContext) ctx, ctx);
        docStarts[0] = 0;
    } else {
        final CompositeReaderContext compCTX = (CompositeReaderContext) ctx;
        final int size = compCTX.leaves().size();
        subSearchers = new SegmentSearcher[size];
        docStarts = new int[size];
        int docBase = 0;
        for (int searcherIDX = 0; searcherIDX < subSearchers.length; searcherIDX++) {
            final LeafReaderContext leave = compCTX.leaves().get(searcherIDX);
            subSearchers[searcherIDX] = new SegmentSearcher(leave, compCTX);
            docStarts[searcherIDX] = docBase;
            docBase += leave.reader().maxDoc();
        }
    }
    final CollapseTopFieldDocs[] shardHits = new CollapseTopFieldDocs[subSearchers.length];
    final Weight weight = searcher.createNormalizedWeight(new MatchAllDocsQuery(), false);
    for (int shardIDX = 0; shardIDX < subSearchers.length; shardIDX++) {
        final SegmentSearcher subSearcher = subSearchers[shardIDX];
        final CollapsingTopDocsCollector c;
        if (numeric) {
            c = CollapsingTopDocsCollector.createNumeric(collapseField.getField(), sort, expectedNumGroups, false);
        } else {
            c = CollapsingTopDocsCollector.createKeyword(collapseField.getField(), sort, expectedNumGroups, false);
        }
        subSearcher.search(weight, c);
        shardHits[shardIDX] = c.getTopDocs();
    }
    CollapseTopFieldDocs mergedFieldDocs = CollapseTopFieldDocs.merge(sort, 0, expectedNumGroups, shardHits);
    assertTopDocsEquals(mergedFieldDocs, collapseTopFieldDocs);
    w.close();
    reader.close();
    dir.close();
}

Example 3 with Weight

use of org.apache.lucene.search.Weight in project elasticsearch by elastic.

the class BestBucketsDeferringCollector method prepareSelectedBuckets.

/**
     * Replay the wrapped collector, but only on a selection of buckets.
     */
@Override
public void prepareSelectedBuckets(long... selectedBuckets) throws IOException {
    if (!finished) {
        throw new IllegalStateException("Cannot replay yet, collection is not finished: postCollect() has not been called");
    }
    if (this.selectedBuckets != null) {
        throw new IllegalStateException("Already been replayed");
    }
    final LongHash hash = new LongHash(selectedBuckets.length, BigArrays.NON_RECYCLING_INSTANCE);
    for (long bucket : selectedBuckets) {
        hash.add(bucket);
    }
    this.selectedBuckets = hash;
    boolean needsScores = collector.needsScores();
    Weight weight = null;
    if (needsScores) {
        weight = searchContext.searcher().createNormalizedWeight(searchContext.query(), true);
    }
    for (Entry entry : entries) {
        final LeafBucketCollector leafCollector = collector.getLeafCollector(entry.context);
        DocIdSetIterator docIt = null;
        if (needsScores && entry.docDeltas.size() > 0) {
            Scorer scorer = weight.scorer(entry.context);
            // We don't need to check if the scorer is null
            // since we are sure that there are documents to replay (entry.docDeltas it not empty).
            docIt = scorer.iterator();
            leafCollector.setScorer(scorer);
        }
        final PackedLongValues.Iterator docDeltaIterator = entry.docDeltas.iterator();
        final PackedLongValues.Iterator buckets = entry.buckets.iterator();
        int doc = 0;
        for (long i = 0, end = entry.docDeltas.size(); i < end; ++i) {
            doc += docDeltaIterator.next();
            final long bucket = buckets.next();
            final long rebasedBucket = hash.find(bucket);
            if (rebasedBucket != -1) {
                if (needsScores) {
                    if (docIt.docID() < doc) {
                        docIt.advance(doc);
                    }
                    // aggregations should only be replayed on matching documents
                    assert docIt.docID() == doc;
                }
                leafCollector.collect(doc, rebasedBucket);
            }
        }
    }
    collector.postCollection();
}

Example 4 with Weight

use of org.apache.lucene.search.Weight in project elasticsearch by elastic.

the class PercolateQuery method createWeight.

@Override
public Weight createWeight(IndexSearcher searcher, boolean needsScores) throws IOException {
    final Weight verifiedMatchesWeight = verifiedMatchesQuery.createWeight(searcher, false);
    final Weight candidateMatchesWeight = candidateMatchesQuery.createWeight(searcher, false);
    return new Weight(this) {

        @Override
        public void extractTerms(Set<Term> set) {
        }

        @Override
        public Explanation explain(LeafReaderContext leafReaderContext, int docId) throws IOException {
            Scorer scorer = scorer(leafReaderContext);
            if (scorer != null) {
                TwoPhaseIterator twoPhaseIterator = scorer.twoPhaseIterator();
                int result = twoPhaseIterator.approximation().advance(docId);
                if (result == docId) {
                    if (twoPhaseIterator.matches()) {
                        if (needsScores) {
                            CheckedFunction<Integer, Query, IOException> percolatorQueries = queryStore.getQueries(leafReaderContext);
                            Query query = percolatorQueries.apply(docId);
                            Explanation detail = percolatorIndexSearcher.explain(query, 0);
                            return Explanation.match(scorer.score(), "PercolateQuery", detail);
                        } else {
                            return Explanation.match(scorer.score(), "PercolateQuery");
                        }
                    }
                }
            }
            return Explanation.noMatch("PercolateQuery");
        }

        @Override
        public float getValueForNormalization() throws IOException {
            return candidateMatchesWeight.getValueForNormalization();
        }

        @Override
        public void normalize(float v, float v1) {
            candidateMatchesWeight.normalize(v, v1);
        }

        @Override
        public Scorer scorer(LeafReaderContext leafReaderContext) throws IOException {
            final Scorer approximation = candidateMatchesWeight.scorer(leafReaderContext);
            if (approximation == null) {
                return null;
            }
            final CheckedFunction<Integer, Query, IOException> queries = queryStore.getQueries(leafReaderContext);
            if (needsScores) {
                return new BaseScorer(this, approximation, queries, percolatorIndexSearcher) {

                    float score;

                    @Override
                    boolean matchDocId(int docId) throws IOException {
                        Query query = percolatorQueries.apply(docId);
                        if (query != null) {
                            TopDocs topDocs = percolatorIndexSearcher.search(query, 1);
                            if (topDocs.totalHits > 0) {
                                score = topDocs.scoreDocs[0].score;
                                return true;
                            } else {
                                return false;
                            }
                        } else {
                            return false;
                        }
                    }

                    @Override
                    public float score() throws IOException {
                        return score;
                    }
                };
            } else {
                Scorer verifiedDocsScorer = verifiedMatchesWeight.scorer(leafReaderContext);
                Bits verifiedDocsBits = Lucene.asSequentialAccessBits(leafReaderContext.reader().maxDoc(), verifiedDocsScorer);
                return new BaseScorer(this, approximation, queries, percolatorIndexSearcher) {

                    @Override
                    public float score() throws IOException {
                        return 0f;
                    }

                    boolean matchDocId(int docId) throws IOException {
                        // the MemoryIndex verification.
                        if (verifiedDocsBits.get(docId)) {
                            return true;
                        }
                        Query query = percolatorQueries.apply(docId);
                        return query != null && Lucene.exists(percolatorIndexSearcher, query);
                    }
                };
            }
        }
    };
}

Example 5 with Weight

use of org.apache.lucene.search.Weight in project elasticsearch by elastic.

the class FunctionScoreTests method testPropagatesApproximations.

public void testPropagatesApproximations() throws IOException {
    Query query = new RandomApproximationQuery(new MatchAllDocsQuery(), random());
    IndexSearcher searcher = newSearcher(reader);
    // otherwise we could get a cached entry that does not have approximations
    searcher.setQueryCache(null);
    FunctionScoreQuery fsq = new FunctionScoreQuery(query, null, null, null, Float.POSITIVE_INFINITY);
    for (boolean needsScores : new boolean[] { true, false }) {
        Weight weight = searcher.createWeight(fsq, needsScores);
        Scorer scorer = weight.scorer(reader.leaves().get(0));
        assertNotNull(scorer.twoPhaseIterator());
    }
    FiltersFunctionScoreQuery ffsq = new FiltersFunctionScoreQuery(query, ScoreMode.SUM, new FilterFunction[0], Float.POSITIVE_INFINITY, null, CombineFunction.MULTIPLY);
    for (boolean needsScores : new boolean[] { true, false }) {
        Weight weight = searcher.createWeight(ffsq, needsScores);
        Scorer scorer = weight.scorer(reader.leaves().get(0));
        assertNotNull(scorer.twoPhaseIterator());
    }
}