Example usage for org.apache.lucene.util DocIdSetBuilder grow

List of usage examples for org.apache.lucene.util DocIdSetBuilder grow

Introduction

In this page you can find the example usage for org.apache.lucene.util DocIdSetBuilder grow.

Prototype

public BulkAdder grow(int numDocs) 

Source Link

Document

Reserve space and return a BulkAdder object that can be used to add up to numDocs documents.

Usage

From source file:org.elasticsearch.search.aggregations.bucket.composite.SortedDocsProducer.java

License:Apache License

/**
 * Visits all non-deleted documents in <code>iterator</code> and fills the provided <code>queue</code>
 * with the top composite buckets extracted from the collection.
 * Documents that contain a top composite bucket are added in the provided <code>builder</code> if it is not null.
 *
 * Returns true if the queue is full and the current <code>leadSourceBucket</code> did not produce any competitive
 * composite buckets.//  w  ww  .  jav a2 s.c  o m
 */
protected boolean processBucket(CompositeValuesCollectorQueue queue, LeafReaderContext context,
        DocIdSetIterator iterator, Comparable<?> leadSourceBucket, @Nullable DocIdSetBuilder builder)
        throws IOException {
    final int[] topCompositeCollected = new int[1];
    final boolean[] hasCollected = new boolean[1];
    final LeafBucketCollector queueCollector = new LeafBucketCollector() {
        int lastDoc = -1;

        // we need to add the matching document in the builder
        // so we build a bulk adder from the approximate cost of the iterator
        // and rebuild the adder during the collection if needed
        int remainingBits = (int) Math.min(iterator.cost(), Integer.MAX_VALUE);
        DocIdSetBuilder.BulkAdder adder = builder == null ? null : builder.grow(remainingBits);

        @Override
        public void collect(int doc, long bucket) throws IOException {
            hasCollected[0] = true;
            int slot = queue.addIfCompetitive();
            if (slot != -1) {
                topCompositeCollected[0]++;
                if (adder != null && doc != lastDoc) {
                    if (remainingBits == 0) {
                        // the cost approximation was lower than the real size, we need to grow the adder
                        // by some numbers (128) to ensure that we can add the extra documents
                        adder = builder.grow(128);
                        remainingBits = 128;
                    }
                    adder.add(doc);
                    remainingBits--;
                    lastDoc = doc;
                }
            }
        }
    };
    final Bits liveDocs = context.reader().getLiveDocs();
    final LeafBucketCollector collector = queue.getLeafCollector(leadSourceBucket, context, queueCollector);
    while (iterator.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
        if (liveDocs == null || liveDocs.get(iterator.docID())) {
            collector.collect(iterator.docID());
        }
    }
    if (queue.isFull() && hasCollected[0] && topCompositeCollected[0] == 0) {
        return true;
    }
    return false;
}