Example usage for org.apache.lucene.index IndexReader leaves

List of usage examples for org.apache.lucene.index IndexReader leaves

Introduction

In this page you can find the example usage for org.apache.lucene.index IndexReader leaves.

Prototype

public final List<LeafReaderContext> leaves() 

Source Link

Document

Returns the reader's leaves, or itself if this reader is atomic.

Usage

From source file:org.elasticsearch.search.fetch.subphase.ScriptFieldsFetchSubPhase.java

License:Apache License

@Override
public void hitsExecute(SearchContext context, SearchHit[] hits) throws IOException {
    if (context.hasScriptFields() == false) {
        return;//from  ww  w .  j  av  a 2 s.  c  o  m
    }

    hits = hits.clone(); // don't modify the incoming hits
    Arrays.sort(hits, Comparator.comparingInt(SearchHit::docId));

    int lastReaderId = -1;
    SearchScript[] leafScripts = null;
    List<ScriptFieldsContext.ScriptField> scriptFields = context.scriptFields().fields();
    final IndexReader reader = context.searcher().getIndexReader();
    for (SearchHit hit : hits) {
        int readerId = ReaderUtil.subIndex(hit.docId(), reader.leaves());
        LeafReaderContext leafReaderContext = reader.leaves().get(readerId);
        if (readerId != lastReaderId) {
            leafScripts = createLeafScripts(leafReaderContext, scriptFields);
            lastReaderId = readerId;
        }
        int docId = hit.docId() - leafReaderContext.docBase;
        for (int i = 0; i < leafScripts.length; i++) {
            leafScripts[i].setDocument(docId);
            final Object value;
            try {
                value = leafScripts[i].run();
            } catch (RuntimeException e) {
                if (scriptFields.get(i).ignoreException()) {
                    continue;
                }
                throw e;
            }
            if (hit.fieldsOrNull() == null) {
                hit.fields(new HashMap<>(2));
            }
            String scriptFieldName = scriptFields.get(i).name();
            DocumentField hitField = hit.getFields().get(scriptFieldName);
            if (hitField == null) {
                final List<Object> values;
                if (value instanceof Collection) {
                    values = new ArrayList<>((Collection<?>) value);
                } else {
                    values = Collections.singletonList(value);
                }
                hitField = new DocumentField(scriptFieldName, values);
                hit.getFields().put(scriptFieldName, hitField);
            }
        }
    }
}

From source file:org.elasticsearch.search.query.TopDocsCollectorContext.java

License:Apache License

/**
 * Returns query total hit count if the <code>query</code> is a {@link MatchAllDocsQuery}
 * or a {@link TermQuery} and the <code>reader</code> has no deletions,
 * -1 otherwise./*from  w  ww .  jav a 2s .  c o  m*/
 */
static int shortcutTotalHitCount(IndexReader reader, Query query) throws IOException {
    while (true) {
        // remove wrappers that don't matter for counts
        // this is necessary so that we don't only optimize match_all
        // queries but also match_all queries that are nested in
        // a constant_score query
        if (query instanceof ConstantScoreQuery) {
            query = ((ConstantScoreQuery) query).getQuery();
        } else if (query instanceof BoostQuery) {
            query = ((BoostQuery) query).getQuery();
        } else {
            break;
        }
    }
    if (query.getClass() == MatchAllDocsQuery.class) {
        return reader.numDocs();
    } else if (query.getClass() == TermQuery.class && reader.hasDeletions() == false) {
        final Term term = ((TermQuery) query).getTerm();
        int count = 0;
        for (LeafReaderContext context : reader.leaves()) {
            count += context.reader().docFreq(term);
        }
        return count;
    } else {
        return -1;
    }
}

From source file:org.elasticsearch.search.suggest.completion.Completion090PostingsFormat.java

License:Apache License

public CompletionStats completionStats(IndexReader indexReader, String... fields) {
    CompletionStats completionStats = new CompletionStats();
    for (AtomicReaderContext atomicReaderContext : indexReader.leaves()) {
        AtomicReader atomicReader = atomicReaderContext.reader();
        try {/*w  ww  .  j a va2s .c  om*/
            for (String fieldName : atomicReader.fields()) {
                Terms terms = atomicReader.fields().terms(fieldName);
                if (terms instanceof CompletionTerms) {
                    CompletionTerms completionTerms = (CompletionTerms) terms;
                    completionStats.add(completionTerms.stats(fields));
                }
            }
        } catch (IOException e) {
            logger.error("Could not get completion stats: {}", e, e.getMessage());
        }
    }

    return completionStats;
}

From source file:org.elasticsearch.search.suggest.completion.CompletionFieldStats.java

License:Apache License

/**
 * Returns total in-heap bytes used by all suggesters.  This method has CPU cost <code>O(numIndexedFields)</code>.
 *
 * @param fieldNamePatterns if non-null, any completion field name matching any of these patterns will break out its in-heap bytes
 * separately in the returned {@link CompletionStats}
 */// ww  w. j a v a 2 s  . c o  m
public static CompletionStats completionStats(IndexReader indexReader, String... fieldNamePatterns) {
    long sizeInBytes = 0;
    ObjectLongHashMap<String> completionFields = null;
    if (fieldNamePatterns != null && fieldNamePatterns.length > 0) {
        completionFields = new ObjectLongHashMap<>(fieldNamePatterns.length);
    }
    for (LeafReaderContext atomicReaderContext : indexReader.leaves()) {
        LeafReader atomicReader = atomicReaderContext.reader();
        try {
            Fields fields = atomicReader.fields();
            for (String fieldName : fields) {
                Terms terms = fields.terms(fieldName);
                if (terms instanceof CompletionTerms) {
                    // TODO: currently we load up the suggester for reporting its size
                    long fstSize = ((CompletionTerms) terms).suggester().ramBytesUsed();
                    if (fieldNamePatterns != null && fieldNamePatterns.length > 0
                            && Regex.simpleMatch(fieldNamePatterns, fieldName)) {
                        completionFields.addTo(fieldName, fstSize);
                    }
                    sizeInBytes += fstSize;
                }
            }
        } catch (IOException ioe) {
            throw new ElasticsearchException(ioe);
        }
    }
    return new CompletionStats(sizeInBytes, completionFields);
}

From source file:org.elasticsearch.search.suggest.completion.CompletionSuggester.java

License:Apache License

@Override
protected Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>> innerExecute(
        String name, CompletionSuggestionContext suggestionContext, IndexReader indexReader, CharsRef spare)
        throws IOException {
    if (suggestionContext.mapper() == null || !(suggestionContext.mapper() instanceof CompletionFieldMapper)) {
        throw new ElasticsearchException(
                "Field [" + suggestionContext.getField() + "] is not a completion suggest field");
    }//from w w w .java  2s  .  c om

    CompletionSuggestion completionSuggestion = new CompletionSuggestion(name, suggestionContext.getSize());
    UnicodeUtil.UTF8toUTF16(suggestionContext.getText(), spare);

    CompletionSuggestion.Entry completionSuggestEntry = new CompletionSuggestion.Entry(
            new StringText(spare.toString()), 0, spare.length());
    completionSuggestion.addTerm(completionSuggestEntry);

    String fieldName = suggestionContext.getField();
    Map<String, CompletionSuggestion.Entry.Option> results = Maps
            .newHashMapWithExpectedSize(indexReader.leaves().size() * suggestionContext.getSize());
    for (AtomicReaderContext atomicReaderContext : indexReader.leaves()) {
        AtomicReader atomicReader = atomicReaderContext.reader();
        Terms terms = atomicReader.fields().terms(fieldName);
        if (terms instanceof Completion090PostingsFormat.CompletionTerms) {
            final Completion090PostingsFormat.CompletionTerms lookupTerms = (Completion090PostingsFormat.CompletionTerms) terms;
            final Lookup lookup = lookupTerms.getLookup(suggestionContext.mapper(), suggestionContext);
            if (lookup == null) {
                // we don't have a lookup for this segment.. this might be possible if a merge dropped all
                // docs from the segment that had a value in this segment.
                continue;
            }
            List<Lookup.LookupResult> lookupResults = lookup.lookup(spare, false, suggestionContext.getSize());
            for (Lookup.LookupResult res : lookupResults) {

                final String key = res.key.toString();
                final float score = res.value;
                final Option value = results.get(key);
                if (value == null) {
                    final Option option = new CompletionSuggestion.Entry.Option(new StringText(key), score,
                            res.payload == null ? null : new BytesArray(res.payload));
                    results.put(key, option);
                } else if (value.getScore() < score) {
                    value.setScore(score);
                    value.setPayload(res.payload == null ? null : new BytesArray(res.payload));
                }
            }
        }
    }
    final List<CompletionSuggestion.Entry.Option> options = new ArrayList<CompletionSuggestion.Entry.Option>(
            results.values());
    CollectionUtil.introSort(options, scoreComparator);

    int optionCount = Math.min(suggestionContext.getSize(), options.size());
    for (int i = 0; i < optionCount; i++) {
        completionSuggestEntry.addOption(options.get(i));
    }

    return completionSuggestion;
}

From source file:org.elasticsearch.search.suggest.completion.old.Completion090PostingsFormat.java

License:Apache License

public CompletionStats completionStats(IndexReader indexReader, String... fields) {
    CompletionStats completionStats = new CompletionStats();
    for (LeafReaderContext atomicReaderContext : indexReader.leaves()) {
        LeafReader atomicReader = atomicReaderContext.reader();
        try {// w  w  w.  j a va2s. c  o m
            for (String fieldName : atomicReader.fields()) {
                Terms terms = atomicReader.fields().terms(fieldName);
                if (terms instanceof CompletionTerms) {
                    CompletionTerms completionTerms = (CompletionTerms) terms;
                    completionStats.add(completionTerms.stats(fields));
                }
            }
        } catch (IOException e) {
            logger.error("Could not get completion stats: {}", e, e.getMessage());
        }
    }

    return completionStats;
}

From source file:org.elasticsearch.search.suggest.completion.old.CompletionSuggester.java

License:Apache License

@Override
protected Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>> innerExecute(
        String name, CompletionSuggestionContext suggestionContext, IndexSearcher searcher,
        CharsRefBuilder spare) throws IOException {
    if (suggestionContext.fieldType() == null) {
        throw new ElasticsearchException(
                "Field [" + suggestionContext.getField() + "] is not a completion suggest field");
    }/*from   ww w .  jav  a2s .  co m*/
    final IndexReader indexReader = searcher.getIndexReader();
    CompletionSuggestion completionSuggestion = new CompletionSuggestion(name, suggestionContext.getSize());
    spare.copyUTF8Bytes(suggestionContext.getText());

    CompletionSuggestion.Entry completionSuggestEntry = new CompletionSuggestion.Entry(
            new StringText(spare.toString()), 0, spare.length());
    completionSuggestion.addTerm(completionSuggestEntry);

    String fieldName = suggestionContext.getField();
    Map<String, CompletionSuggestion.Entry.Option> results = Maps
            .newHashMapWithExpectedSize(indexReader.leaves().size() * suggestionContext.getSize());
    for (LeafReaderContext atomicReaderContext : indexReader.leaves()) {
        LeafReader atomicReader = atomicReaderContext.reader();
        Terms terms = atomicReader.fields().terms(fieldName);
        if (terms instanceof Completion090PostingsFormat.CompletionTerms) {
            final Completion090PostingsFormat.CompletionTerms lookupTerms = (Completion090PostingsFormat.CompletionTerms) terms;
            final Lookup lookup = lookupTerms.getLookup(suggestionContext.fieldType(), suggestionContext);
            if (lookup == null) {
                // we don't have a lookup for this segment.. this might be possible if a merge dropped all
                // docs from the segment that had a value in this segment.
                continue;
            }
            List<Lookup.LookupResult> lookupResults = lookup.lookup(spare.get(), false,
                    suggestionContext.getSize());
            for (Lookup.LookupResult res : lookupResults) {

                final String key = res.key.toString();
                final float score = res.value;
                final CompletionSuggestion.Entry.Option value = results.get(key);
                if (value == null) {
                    final CompletionSuggestion.Entry.Option option = new CompletionSuggestion.Entry.Option(
                            new StringText(key), score,
                            res.payload == null ? null : new BytesArray(res.payload));
                    results.put(key, option);
                } else if (value.getScore() < score) {
                    value.setScore(score);
                    value.setPayload(res.payload == null ? null : new BytesArray(res.payload));
                }
            }
        }
    }
    final List<CompletionSuggestion.Entry.Option> options = new ArrayList<>(results.values());
    CollectionUtil.introSort(options, scoreComparator);

    int optionCount = Math.min(suggestionContext.getSize(), options.size());
    for (int i = 0; i < optionCount; i++) {
        completionSuggestEntry.addOption(options.get(i));
    }

    return completionSuggestion;
}

From source file:org.exist.indexing.lucene.LuceneIndexWorker.java

License:Open Source License

public String getFieldContent(int docId, String field) throws IOException {
    BytesRef bytes = new BytesRef(NumericUtils.BUF_SIZE_INT);
    NumericUtils.intToPrefixCoded(docId, 0, bytes);
    Term dt = new Term(FIELD_DOC_ID, bytes);

    IndexReader reader = null;
    try {/*  w  w w.j  a v a  2s.  c  o  m*/
        reader = index.getReader();

        List<AtomicReaderContext> leaves = reader.leaves();
        for (AtomicReaderContext context : leaves) {
            DocsEnum docs = context.reader().termDocsEnum(dt);
            if (docs != null && docs.nextDoc() != DocsEnum.NO_MORE_DOCS) {
                Document doc = reader.document(docs.docID());
                return doc.get(field);
            }
        }
    } finally {
        index.releaseReader(reader);
    }
    return null;
}

From source file:org.exist.indexing.lucene.LuceneIndexWorker.java

License:Open Source License

public boolean hasIndex(int docId) throws IOException {
    BytesRef bytes = new BytesRef(NumericUtils.BUF_SIZE_INT);
    NumericUtils.intToPrefixCoded(docId, 0, bytes);
    Term dt = new Term(FIELD_DOC_ID, bytes);

    IndexReader reader = null;
    try {/* w  w  w .j  a v a2s . com*/
        reader = index.getReader();

        boolean found = false;
        List<AtomicReaderContext> leaves = reader.leaves();
        for (AtomicReaderContext context : leaves) {
            DocsEnum docs = context.reader().termDocsEnum(dt);
            if (docs != null && docs.nextDoc() != DocsEnum.NO_MORE_DOCS) {
                found = true;
                break;
            }
        }
        return found;
    } finally {
        index.releaseReader(reader);
    }
}

From source file:org.exist.indexing.lucene.LuceneIndexWorker.java

License:Open Source License

private Occurrences[] scanIndexByQName(List<QName> qnames, DocumentSet docs, NodeSet nodes, String start,
        String end, long max) {
    TreeMap<String, Occurrences> map = new TreeMap<>();
    IndexReader reader = null;
    try {/*from w  ww.  j  a  v  a2s.  c o m*/
        reader = index.getReader();
        for (QName qname : qnames) {
            String field = LuceneUtil.encodeQName(qname, index.getBrokerPool().getSymbols());
            List<AtomicReaderContext> leaves = reader.leaves();
            for (AtomicReaderContext context : leaves) {
                NumericDocValues docIdValues = context.reader().getNumericDocValues(FIELD_DOC_ID);
                BinaryDocValues nodeIdValues = context.reader().getBinaryDocValues(LuceneUtil.FIELD_NODE_ID);
                Bits liveDocs = context.reader().getLiveDocs();
                Terms terms = context.reader().terms(field);
                if (terms == null)
                    continue;
                TermsEnum termsIter = terms.iterator(null);
                if (termsIter.next() == null) {
                    continue;
                }
                do {
                    if (map.size() >= max) {
                        break;
                    }
                    BytesRef ref = termsIter.term();
                    String term = ref.utf8ToString();
                    boolean include = true;
                    if (end != null) {
                        if (term.compareTo(end) > 0)
                            include = false;
                    } else if (start != null && !term.startsWith(start))
                        include = false;
                    if (include) {
                        DocsEnum docsEnum = termsIter.docs(null, null);
                        while (docsEnum.nextDoc() != DocsEnum.NO_MORE_DOCS) {
                            if (liveDocs != null && !liveDocs.get(docsEnum.docID())) {
                                continue;
                            }
                            int docId = (int) docIdValues.get(docsEnum.docID());
                            DocumentImpl storedDocument = docs.getDoc(docId);
                            if (storedDocument == null)
                                continue;
                            NodeId nodeId = null;
                            if (nodes != null) {
                                BytesRef nodeIdRef = new BytesRef(buf);
                                nodeIdValues.get(docsEnum.docID(), nodeIdRef);
                                int units = ByteConversion.byteToShort(nodeIdRef.bytes, nodeIdRef.offset);
                                nodeId = index.getBrokerPool().getNodeFactory().createFromData(units,
                                        nodeIdRef.bytes, nodeIdRef.offset + 2);
                            }
                            // DW: warning: nodes can be null?
                            if (nodeId == null || nodes.get(storedDocument, nodeId) != null) {
                                Occurrences oc = map.get(term);
                                if (oc == null) {
                                    oc = new Occurrences(term);
                                    map.put(term, oc);
                                }
                                oc.addDocument(storedDocument);
                                oc.addOccurrences(docsEnum.freq());
                            }
                        }
                    }
                } while (termsIter.next() != null);
            }
        }
    } catch (IOException e) {
        LOG.warn("Error while scanning lucene index entries: " + e.getMessage(), e);
    } finally {
        index.releaseReader(reader);
    }
    Occurrences[] occur = new Occurrences[map.size()];
    return map.values().toArray(occur);
}