/*
 * 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.util;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.apache.lucene.index.PointValues;
import org.apache.lucene.index.Terms;
import org.apache.lucene.search.DocIdSet;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.packed.PackedInts;

A builder of DocIdSets. At first it uses a sparse structure to gather documents, and then upgrades to a non-sparse bit set once enough hits match. To add documents, you first need to call grow in order to reserve space, and then call BulkAdder.add(int) on the returned BulkAdder.
@lucene.internal
/** * A builder of {@link DocIdSet}s. At first it uses a sparse structure to gather * documents, and then upgrades to a non-sparse bit set once enough hits match. * * To add documents, you first need to call {@link #grow} in order to reserve * space, and then call {@link BulkAdder#add(int)} on the returned * {@link BulkAdder}. * * @lucene.internal */
public final class DocIdSetBuilder {
Utility class to efficiently add many docs in one go. @see DocIdSetBuilder#grow
/** Utility class to efficiently add many docs in one go. * @see DocIdSetBuilder#grow */
public static abstract class BulkAdder { public abstract void add(int doc); } private static class FixedBitSetAdder extends BulkAdder { final FixedBitSet bitSet; FixedBitSetAdder(FixedBitSet bitSet) { this.bitSet = bitSet; } @Override public void add(int doc) { bitSet.set(doc); } } private static class Buffer { int[] array; int length; Buffer(int length) { this.array = new int[length]; this.length = 0; } Buffer(int[] array, int length) { this.array = array; this.length = length; } } private static class BufferAdder extends BulkAdder { final Buffer buffer; BufferAdder(Buffer buffer) { this.buffer = buffer; } @Override public void add(int doc) { buffer.array[buffer.length++] = doc; } } private final int maxDoc; private final int threshold; // pkg-private for testing final boolean multivalued; final double numValuesPerDoc; private List<Buffer> buffers = new ArrayList<>(); private int totalAllocated; // accumulated size of the allocated buffers private FixedBitSet bitSet; private long counter = -1; private BulkAdder adder;
Create a builder that can contain doc IDs between 0 and maxDoc.
/** * Create a builder that can contain doc IDs between {@code 0} and {@code maxDoc}. */
public DocIdSetBuilder(int maxDoc) { this(maxDoc, -1, -1); }
Create a DocIdSetBuilder instance that is optimized for accumulating docs that match the given Terms.
/** Create a {@link DocIdSetBuilder} instance that is optimized for * accumulating docs that match the given {@link Terms}. */
public DocIdSetBuilder(int maxDoc, Terms terms) throws IOException { this(maxDoc, terms.getDocCount(), terms.getSumDocFreq()); }
Create a DocIdSetBuilder instance that is optimized for accumulating docs that match the given PointValues.
/** Create a {@link DocIdSetBuilder} instance that is optimized for * accumulating docs that match the given {@link PointValues}. */
public DocIdSetBuilder(int maxDoc, PointValues values, String field) throws IOException { this(maxDoc, values.getDocCount(), values.size()); } DocIdSetBuilder(int maxDoc, int docCount, long valueCount) { this.maxDoc = maxDoc; this.multivalued = docCount < 0 || docCount != valueCount; if (docCount <= 0 || valueCount < 0) { // assume one value per doc, this means the cost will be overestimated // if the docs are actually multi-valued this.numValuesPerDoc = 1; } else { // otherwise compute from index stats this.numValuesPerDoc = (double) valueCount / docCount; } assert numValuesPerDoc >= 1: "valueCount=" + valueCount + " docCount=" + docCount; // For ridiculously small sets, we'll just use a sorted int[] // maxDoc >>> 7 is a good value if you want to save memory, lower values // such as maxDoc >>> 11 should provide faster building but at the expense // of using a full bitset even for quite sparse data this.threshold = maxDoc >>> 7; this.bitSet = null; }
Add the content of the provided DocIdSetIterator to this builder. NOTE: if you need to build a DocIdSet out of a single DocIdSetIterator, you should rather use Builder.
/** * Add the content of the provided {@link DocIdSetIterator} to this builder. * NOTE: if you need to build a {@link DocIdSet} out of a single * {@link DocIdSetIterator}, you should rather use {@link RoaringDocIdSet.Builder}. */
public void add(DocIdSetIterator iter) throws IOException { if (bitSet != null) { bitSet.or(iter); return; } int cost = (int) Math.min(Integer.MAX_VALUE, iter.cost()); BulkAdder adder = grow(cost); for (int i = 0; i < cost; ++i) { int doc = iter.nextDoc(); if (doc == DocIdSetIterator.NO_MORE_DOCS) { return; } adder.add(doc); } for (int doc = iter.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = iter.nextDoc()) { grow(1).add(doc); } }
Reserve space and return a BulkAdder object that can be used to add up to numDocs documents.
/** * Reserve space and return a {@link BulkAdder} object that can be used to * add up to {@code numDocs} documents. */
public BulkAdder grow(int numDocs) { if (bitSet == null) { if ((long) totalAllocated + numDocs <= threshold) { ensureBufferCapacity(numDocs); } else { upgradeToBitSet(); counter += numDocs; } } else { counter += numDocs; } return adder; } private void ensureBufferCapacity(int numDocs) { if (buffers.isEmpty()) { addBuffer(additionalCapacity(numDocs)); return; } Buffer current = buffers.get(buffers.size() - 1); if (current.array.length - current.length >= numDocs) { // current buffer is large enough return; } if (current.length < current.array.length - (current.array.length >>> 3)) { // current buffer is less than 7/8 full, resize rather than waste space growBuffer(current, additionalCapacity(numDocs)); } else { addBuffer(additionalCapacity(numDocs)); } } private int additionalCapacity(int numDocs) { // exponential growth: the new array has a size equal to the sum of what // has been allocated so far int c = totalAllocated; // but is also >= numDocs + 1 so that we can store the next batch of docs // (plus an empty slot so that we are more likely to reuse the array in build()) c = Math.max(numDocs + 1, c); // avoid cold starts c = Math.max(32, c); // do not go beyond the threshold c = Math.min(threshold - totalAllocated, c); return c; } private Buffer addBuffer(int len) { Buffer buffer = new Buffer(len); buffers.add(buffer); adder = new BufferAdder(buffer); totalAllocated += buffer.array.length; return buffer; } private void growBuffer(Buffer buffer, int additionalCapacity) { buffer.array = ArrayUtil.growExact(buffer.array, buffer.array.length + additionalCapacity); totalAllocated += additionalCapacity; } private void upgradeToBitSet() { assert bitSet == null; FixedBitSet bitSet = new FixedBitSet(maxDoc); long counter = 0; for (Buffer buffer : buffers) { int[] array = buffer.array; int length = buffer.length; counter += length; for (int i = 0; i < length; ++i) { bitSet.set(array[i]); } } this.bitSet = bitSet; this.counter = counter; this.buffers = null; this.adder = new FixedBitSetAdder(bitSet); }
Build a DocIdSet from the accumulated doc IDs.
/** * Build a {@link DocIdSet} from the accumulated doc IDs. */
public DocIdSet build() { try { if (bitSet != null) { assert counter >= 0; final long cost = Math.round(counter / numValuesPerDoc); return new BitDocIdSet(bitSet, cost); } else { Buffer concatenated = concat(buffers); LSBRadixSorter sorter = new LSBRadixSorter(); sorter.sort(PackedInts.bitsRequired(maxDoc - 1), concatenated.array, concatenated.length); final int l; if (multivalued) { l = dedup(concatenated.array, concatenated.length); } else { assert noDups(concatenated.array, concatenated.length); l = concatenated.length; } assert l <= concatenated.length; concatenated.array[l] = DocIdSetIterator.NO_MORE_DOCS; return new IntArrayDocIdSet(concatenated.array, l); } } finally { this.buffers = null; this.bitSet = null; } }
Concatenate the buffers in any order, leaving at least one empty slot in the end NOTE: this method might reuse one of the arrays
/** * Concatenate the buffers in any order, leaving at least one empty slot in * the end * NOTE: this method might reuse one of the arrays */
private static Buffer concat(List<Buffer> buffers) { int totalLength = 0; Buffer largestBuffer = null; for (Buffer buffer : buffers) { totalLength += buffer.length; if (largestBuffer == null || buffer.array.length > largestBuffer.array.length) { largestBuffer = buffer; } } if (largestBuffer == null) { return new Buffer(1); } int[] docs = largestBuffer.array; if (docs.length < totalLength + 1) { docs = ArrayUtil.growExact(docs, totalLength + 1); } totalLength = largestBuffer.length; for (Buffer buffer : buffers) { if (buffer != largestBuffer) { System.arraycopy(buffer.array, 0, docs, totalLength, buffer.length); totalLength += buffer.length; } } return new Buffer(docs, totalLength); } private static int dedup(int[] arr, int length) { if (length == 0) { return 0; } int l = 1; int previous = arr[0]; for (int i = 1; i < length; ++i) { final int value = arr[i]; assert value >= previous; if (value != previous) { arr[l++] = value; previous = value; } } return l; } private static boolean noDups(int[] a, int len) { for (int i = 1; i < len; ++i) { assert a[i-1] < a[i]; } return true; } }