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

Example 6 with LeafCollector

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

the class TermsSliceQueryTests method testSearch.

public void testSearch() throws Exception {
    final int numDocs = randomIntBetween(100, 200);
    final Directory dir = newDirectory();
    final RandomIndexWriter w = new RandomIndexWriter(random(), dir, new KeywordAnalyzer());
    int max = randomIntBetween(2, 10);
    int[] sliceCounters = new int[max];
    Set<String> keys = new HashSet<>();
    for (int i = 0; i < numDocs; ++i) {
        Document doc = new Document();
        String uuid = UUIDs.base64UUID();
        BytesRef br = new BytesRef(uuid);
        int id = Math.floorMod(br.hashCode(), max);
        sliceCounters[id]++;
        doc.add(new StringField("uuid", uuid, Field.Store.YES));
        w.addDocument(doc);
        keys.add(uuid);
    }
    final IndexReader reader = w.getReader();
    final IndexSearcher searcher = newSearcher(reader);
    for (int id = 0; id < max; id++) {
        TermsSliceQuery query1 = new TermsSliceQuery("uuid", id, max);
        assertThat(searcher.count(query1), equalTo(sliceCounters[id]));
        searcher.search(query1, new Collector() {

            @Override
            public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
                return new LeafCollector() {

                    @Override
                    public void setScorer(Scorer scorer) throws IOException {
                    }

                    @Override
                    public void collect(int doc) throws IOException {
                        Document d = context.reader().document(doc, Collections.singleton("uuid"));
                        String uuid = d.get("uuid");
                        assertThat(keys.contains(uuid), equalTo(true));
                        keys.remove(uuid);
                    }
                };
            }

            @Override
            public boolean needsScores() {
                return false;
            }
        });
    }
    assertThat(keys.size(), equalTo(0));
    w.close();
    reader.close();
    dir.close();
}

Example 7 with LeafCollector

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

the class DocValuesSliceQueryTests method testSearch.

public void testSearch() throws Exception {
    final int numDocs = randomIntBetween(100, 200);
    final Directory dir = newDirectory();
    final RandomIndexWriter w = new RandomIndexWriter(random(), dir);
    int max = randomIntBetween(2, 10);
    int[] sliceCounters1 = new int[max];
    int[] sliceCounters2 = new int[max];
    Set<String> keys = new HashSet<>();
    for (int i = 0; i < numDocs; ++i) {
        Document doc = new Document();
        String uuid = UUIDs.base64UUID();
        int intValue = randomInt();
        long doubleValue = NumericUtils.doubleToSortableLong(randomDouble());
        doc.add(new StringField("uuid", uuid, Field.Store.YES));
        doc.add(new SortedNumericDocValuesField("intField", intValue));
        doc.add(new SortedNumericDocValuesField("doubleField", doubleValue));
        w.addDocument(doc);
        sliceCounters1[Math.floorMod(BitMixer.mix((long) intValue), max)]++;
        sliceCounters2[Math.floorMod(BitMixer.mix(doubleValue), max)]++;
        keys.add(uuid);
    }
    final IndexReader reader = w.getReader();
    final IndexSearcher searcher = newSearcher(reader);
    for (int id = 0; id < max; id++) {
        DocValuesSliceQuery query1 = new DocValuesSliceQuery("intField", id, max);
        assertThat(searcher.count(query1), equalTo(sliceCounters1[id]));
        DocValuesSliceQuery query2 = new DocValuesSliceQuery("doubleField", id, max);
        assertThat(searcher.count(query2), equalTo(sliceCounters2[id]));
        searcher.search(query1, new Collector() {

            @Override
            public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
                return new LeafCollector() {

                    @Override
                    public void setScorer(Scorer scorer) throws IOException {
                    }

                    @Override
                    public void collect(int doc) throws IOException {
                        Document d = context.reader().document(doc, Collections.singleton("uuid"));
                        String uuid = d.get("uuid");
                        assertThat(keys.contains(uuid), equalTo(true));
                        keys.remove(uuid);
                    }
                };
            }

            @Override
            public boolean needsScores() {
                return false;
            }
        });
    }
    assertThat(keys.size(), equalTo(0));
    w.close();
    reader.close();
    dir.close();
}

Example 8 with LeafCollector

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

the class SearchCancellationTests method testLowLevelCancellableCollector.

public void testLowLevelCancellableCollector() throws IOException {
    TotalHitCountCollector collector = new TotalHitCountCollector();
    AtomicBoolean cancelled = new AtomicBoolean();
    CancellableCollector cancellableCollector = new CancellableCollector(cancelled::get, true, collector);
    final LeafCollector leafCollector = cancellableCollector.getLeafCollector(reader.leaves().get(0));
    leafCollector.collect(0);
    cancelled.set(true);
    expectThrows(TaskCancelledException.class, () -> leafCollector.collect(1));
}

Example 9 with LeafCollector

use of org.apache.lucene.search.LeafCollector in project neo4j by neo4j.

the class DocValuesCollector method replayTo.

private void replayTo(Collector collector) throws IOException {
    for (MatchingDocs docs : getMatchingDocs()) {
        LeafCollector leafCollector = collector.getLeafCollector(docs.context);
        Scorer scorer;
        DocIdSetIterator idIterator = docs.docIdSet.iterator();
        if (isKeepScores()) {
            scorer = new ReplayingScorer(docs.scores);
        } else {
            scorer = new ConstantScoreScorer(null, Float.NaN, idIterator);
        }
        leafCollector.setScorer(scorer);
        int doc;
        while ((doc = idIterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
            leafCollector.collect(doc);
        }
    }
}

Example 10 with LeafCollector

use of org.apache.lucene.search.LeafCollector in project lucene-solr by apache.

the class TestSort method testSort.

public void testSort() throws Exception {
    Directory dir = new RAMDirectory();
    Field f = new StringField("f", "0", Field.Store.NO);
    Field f2 = new StringField("f2", "0", Field.Store.NO);
    for (int iterCnt = 0; iterCnt < iter; iterCnt++) {
        IndexWriter iw = new IndexWriter(dir, new IndexWriterConfig(new SimpleAnalyzer()).setOpenMode(IndexWriterConfig.OpenMode.CREATE));
        final MyDoc[] mydocs = new MyDoc[ndocs];
        int v1EmptyPercent = 50;
        int v2EmptyPercent = 50;
        int commitCountdown = commitCount;
        for (int i = 0; i < ndocs; i++) {
            MyDoc mydoc = new MyDoc();
            mydoc.doc = i;
            mydocs[i] = mydoc;
            Document document = new Document();
            if (r.nextInt(100) < v1EmptyPercent) {
                mydoc.val = Integer.toString(r.nextInt(maxval));
                f.setStringValue(mydoc.val);
                document.add(f);
            }
            if (r.nextInt(100) < v2EmptyPercent) {
                mydoc.val2 = Integer.toString(r.nextInt(maxval));
                f2.setStringValue(mydoc.val2);
                document.add(f2);
            }
            iw.addDocument(document);
            if (--commitCountdown <= 0) {
                commitCountdown = commitCount;
                iw.commit();
            }
        }
        iw.close();
        Map<String, UninvertingReader.Type> mapping = new HashMap<>();
        mapping.put("f", UninvertingReader.Type.SORTED);
        mapping.put("f2", UninvertingReader.Type.SORTED);
        DirectoryReader reader = UninvertingReader.wrap(DirectoryReader.open(dir), mapping);
        IndexSearcher searcher = new IndexSearcher(reader);
        // System.out.println("segments="+searcher.getIndexReader().getSequentialSubReaders().length);
        assertTrue(reader.leaves().size() > 1);
        for (int i = 0; i < qiter; i++) {
            Filter filt = new Filter() {

                @Override
                public DocIdSet getDocIdSet(LeafReaderContext context, Bits acceptDocs) {
                    return BitsFilteredDocIdSet.wrap(randSet(context.reader().maxDoc()), acceptDocs);
                }

                @Override
                public String toString(String field) {
                    return "TestSortFilter";
                }

                @Override
                public boolean equals(Object other) {
                    return other == this;
                }

                @Override
                public int hashCode() {
                    return System.identityHashCode(this);
                }
            };
            int top = r.nextInt((ndocs >> 3) + 1) + 1;
            final boolean luceneSort = r.nextBoolean();
            final boolean sortMissingLast = !luceneSort && r.nextBoolean();
            final boolean sortMissingFirst = !luceneSort && !sortMissingLast;
            final boolean reverse = r.nextBoolean();
            List<SortField> sfields = new ArrayList<>();
            final boolean secondary = r.nextBoolean();
            final boolean luceneSort2 = r.nextBoolean();
            final boolean sortMissingLast2 = !luceneSort2 && r.nextBoolean();
            final boolean sortMissingFirst2 = !luceneSort2 && !sortMissingLast2;
            final boolean reverse2 = r.nextBoolean();
            if (r.nextBoolean())
                sfields.add(new SortField(null, SortField.Type.SCORE));
            // hit both use-cases of sort-missing-last
            sfields.add(Sorting.getStringSortField("f", reverse, sortMissingLast, sortMissingFirst));
            if (secondary) {
                sfields.add(Sorting.getStringSortField("f2", reverse2, sortMissingLast2, sortMissingFirst2));
            }
            if (r.nextBoolean())
                sfields.add(new SortField(null, SortField.Type.SCORE));
            Sort sort = new Sort(sfields.toArray(new SortField[sfields.size()]));
            final String nullRep = luceneSort || sortMissingFirst && !reverse || sortMissingLast && reverse ? "" : "zzz";
            final String nullRep2 = luceneSort2 || sortMissingFirst2 && !reverse2 || sortMissingLast2 && reverse2 ? "" : "zzz";
            boolean trackScores = r.nextBoolean();
            boolean trackMaxScores = r.nextBoolean();
            boolean scoreInOrder = r.nextBoolean();
            final TopFieldCollector topCollector = TopFieldCollector.create(sort, top, true, trackScores, trackMaxScores);
            final List<MyDoc> collectedDocs = new ArrayList<>();
            // delegate and collect docs ourselves
            Collector myCollector = new FilterCollector(topCollector) {

                @Override
                public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
                    final int docBase = context.docBase;
                    return new FilterLeafCollector(super.getLeafCollector(context)) {

                        @Override
                        public void collect(int doc) throws IOException {
                            super.collect(doc);
                            collectedDocs.add(mydocs[docBase + doc]);
                        }
                    };
                }
            };
            searcher.search(filt, myCollector);
            Collections.sort(collectedDocs, (o1, o2) -> {
                String v1 = o1.val == null ? nullRep : o1.val;
                String v2 = o2.val == null ? nullRep : o2.val;
                int cmp = v1.compareTo(v2);
                if (reverse)
                    cmp = -cmp;
                if (cmp != 0)
                    return cmp;
                if (secondary) {
                    v1 = o1.val2 == null ? nullRep2 : o1.val2;
                    v2 = o2.val2 == null ? nullRep2 : o2.val2;
                    cmp = v1.compareTo(v2);
                    if (reverse2)
                        cmp = -cmp;
                }
                cmp = cmp == 0 ? o1.doc - o2.doc : cmp;
                return cmp;
            });
            TopDocs topDocs = topCollector.topDocs();
            ScoreDoc[] sdocs = topDocs.scoreDocs;
            for (int j = 0; j < sdocs.length; j++) {
                int id = sdocs[j].doc;
                if (id != collectedDocs.get(j).doc) {
                    log.error("Error at pos " + j + "\n\tsortMissingFirst=" + sortMissingFirst + " sortMissingLast=" + sortMissingLast + " reverse=" + reverse + "\n\tEXPECTED=" + collectedDocs);
                }
                assertEquals(id, collectedDocs.get(j).doc);
            }
        }
        reader.close();
    }
    dir.close();
}