--- /dev/null
+package org.apache.lucene.facet.search.sampling;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+import org.apache.lucene.util.PriorityQueue;
+
+import org.apache.lucene.facet.search.ScoredDocIDs;
+import org.apache.lucene.facet.search.ScoredDocIDsIterator;
+import org.apache.lucene.facet.util.ScoredDocIdsUtils;
+
+/**
+ * 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.
+ */
+
+/**
+ * Take random samples of large collections.
+ * @lucene.experimental
+ */
+public class RepeatableSampler extends Sampler {
+
+ private static final Logger logger = Logger.getLogger(RepeatableSampler.class.getName());
+
+ public RepeatableSampler(SamplingParams params) {
+ super(params);
+ }
+
+ @Override
+ protected SampleResult createSample(ScoredDocIDs docids, int actualSize,
+ int sampleSetSize) throws IOException {
+ int[] sampleSet = null;
+ try {
+ sampleSet = repeatableSample(docids, actualSize,
+ sampleSetSize);
+ } catch (IOException e) {
+ if (logger.isLoggable(Level.WARNING)) {
+ logger.log(Level.WARNING, "sampling failed: "+e.getMessage()+" - falling back to no sampling!", e);
+ }
+ return new SampleResult(docids, 1d);
+ }
+
+ ScoredDocIDs sampled = ScoredDocIdsUtils.createScoredDocIDsSubset(docids,
+ sampleSet);
+ if (logger.isLoggable(Level.FINEST)) {
+ logger.finest("******************** " + sampled.size());
+ }
+ return new SampleResult(sampled, sampled.size()/(double)docids.size());
+ }
+
+ /**
+ * Returns <code>sampleSize</code> values from the first <code>collectionSize</code>
+ * locations of <code>collection</code>, chosen using
+ * the <code>TRAVERSAL</code> algorithm. The sample values are not sorted.
+ * @param collection The values from which a sample is wanted.
+ * @param collectionSize The number of values (from the first) from which to draw the sample.
+ * @param sampleSize The number of values to return.
+ * @return An array of values chosen from the collection.
+ * @see Algorithm#TRAVERSAL
+ */
+ private static int[] repeatableSample(ScoredDocIDs collection,
+ int collectionSize, int sampleSize)
+ throws IOException {
+ return repeatableSample(collection, collectionSize,
+ sampleSize, Algorithm.HASHING, Sorted.NO);
+ }
+
+ /**
+ * Returns <code>sampleSize</code> values from the first <code>collectionSize</code>
+ * locations of <code>collection</code>, chosen using <code>algorithm</code>.
+ * @param collection The values from which a sample is wanted.
+ * @param collectionSize The number of values (from the first) from which to draw the sample.
+ * @param sampleSize The number of values to return.
+ * @param algorithm Which algorithm to use.
+ * @param sorted Sorted.YES to sort the sample values in ascending order before returning;
+ * Sorted.NO to return them in essentially random order.
+ * @return An array of values chosen from the collection.
+ */
+ private static int[] repeatableSample(ScoredDocIDs collection,
+ int collectionSize, int sampleSize,
+ Algorithm algorithm, Sorted sorted)
+ throws IOException {
+ if (collection == null) {
+ throw new IOException("docIdSet is null");
+ }
+ if (sampleSize < 1) {
+ throw new IOException("sampleSize < 1 (" + sampleSize + ")");
+ }
+ if (collectionSize < sampleSize) {
+ throw new IOException("collectionSize (" + collectionSize + ") less than sampleSize (" + sampleSize + ")");
+ }
+ int[] sample = new int[sampleSize];
+ long[] times = new long[4];
+ if (algorithm == Algorithm.TRAVERSAL) {
+ sample1(collection, collectionSize, sample, times);
+ } else if (algorithm == Algorithm.HASHING) {
+ sample2(collection, collectionSize, sample, times);
+ } else {
+ throw new IllegalArgumentException("Invalid algorithm selection");
+ }
+ if (sorted == Sorted.YES) {
+ Arrays.sort(sample);
+ }
+ if (returnTimings) {
+ times[3] = System.currentTimeMillis();
+ if (logger.isLoggable(Level.FINEST)) {
+ logger.finest("Times: " + (times[1] - times[0]) + "ms, "
+ + (times[2] - times[1]) + "ms, " + (times[3] - times[2])+"ms");
+ }
+ }
+ return sample;
+ }
+
+ /**
+ * Returns <code>sample</code>.length values chosen from the first <code>collectionSize</code>
+ * locations of <code>collection</code>, using the TRAVERSAL algorithm. The sample is
+ * pseudorandom: no subset of the original collection
+ * is in principle more likely to occur than any other, but for a given collection
+ * and sample size, the same sample will always be returned. This algorithm walks the
+ * original collection in a methodical way that is guaranteed not to visit any location
+ * more than once, which makes sampling without replacement faster because removals don't
+ * have to be tracked, and the number of operations is proportional to the sample size,
+ * not the collection size.
+ * Times for performance measurement
+ * are returned in <code>times</code>, which must be an array of at least three longs, containing
+ * nanosecond event times. The first
+ * is set when the algorithm starts; the second, when the step size has been calculated;
+ * and the third when the sample has been taken.
+ * @param collection The set to be sampled.
+ * @param collectionSize The number of values to use (starting from first).
+ * @param sample The array in which to return the sample.
+ * @param times The times of three events, for measuring performance.
+ */
+ private static void sample1(ScoredDocIDs collection, int collectionSize, int[] sample, long[] times)
+ throws IOException {
+ ScoredDocIDsIterator it = collection.iterator();
+ if (returnTimings) {
+ times[0] = System.currentTimeMillis();
+ }
+ int sampleSize = sample.length;
+ int prime = findGoodStepSize(collectionSize, sampleSize);
+ int mod = prime % collectionSize;
+ if (returnTimings) {
+ times[1] = System.currentTimeMillis();
+ }
+ int sampleCount = 0;
+ int index = 0;
+ for (; sampleCount < sampleSize;) {
+ if (index + mod < collectionSize) {
+ for (int i = 0; i < mod; i++, index++) {
+ it.next();
+ }
+ } else {
+ index = index + mod - collectionSize;
+ it = collection.iterator();
+ for (int i = 0; i < index; i++) {
+ it.next();
+ }
+ }
+ sample[sampleCount++] = it.getDocID();
+ }
+ if (returnTimings) {
+ times[2] = System.currentTimeMillis();
+ }
+ }
+
+ /**
+ * Returns a value which will allow the caller to walk
+ * a collection of <code>collectionSize</code> values, without repeating or missing
+ * any, and spanning the collection from beginning to end at least once with
+ * <code>sampleSize</code> visited locations. Choosing a value
+ * that is relatively prime to the collection size ensures that stepping by that size (modulo
+ * the collection size) will hit all locations without repeating, eliminating the need to
+ * track previously visited locations for a "without replacement" sample. Starting with the
+ * square root of the collection size ensures that either the first or second prime tried will
+ * work (they can't both divide the collection size). It also has the property that N steps of
+ * size N will span a collection of N**2 elements once. If the sample is bigger than N, it will
+ * wrap multiple times (without repeating). If the sample is smaller, a step size is chosen
+ * that will result in at least one spanning of the collection.
+ *
+ * @param collectionSize The number of values in the collection to be sampled.
+ * @param sampleSize The number of values wanted in the sample.
+ * @return A good increment value for walking the collection.
+ */
+ private static int findGoodStepSize(int collectionSize, int sampleSize) {
+ int i = (int) Math.sqrt(collectionSize);
+ if (sampleSize < i) {
+ i = collectionSize / sampleSize;
+ }
+ do {
+ i = findNextPrimeAfter(i);
+ } while (collectionSize % i == 0);
+ return i;
+ }
+
+ /**
+ * Returns the first prime number that is larger than <code>n</code>.
+ * @param n A number less than the prime to be returned.
+ * @return The smallest prime larger than <code>n</code>.
+ */
+ private static int findNextPrimeAfter(int n) {
+ n += (n % 2 == 0) ? 1 : 2; // next odd
+ foundFactor: for (;; n += 2) { //TODO labels??!!
+ int sri = (int) (Math.sqrt(n));
+ for (int primeIndex = 0; primeIndex < N_PRIMES; primeIndex++) {
+ int p = primes[primeIndex];
+ if (p > sri) {
+ return n;
+ }
+ if (n % p == 0) {
+ continue foundFactor;
+ }
+ }
+ for (int p = primes[N_PRIMES - 1] + 2;; p += 2) {
+ if (p > sri) {
+ return n;
+ }
+ if (n % p == 0) {
+ continue foundFactor;
+ }
+ }
+ }
+ }
+
+ /**
+ * The first N_PRIMES primes, after 2.
+ */
+ private static final int N_PRIMES = 4000;
+ private static int[] primes = new int[N_PRIMES];
+ static {
+ primes[0] = 3;
+ for (int count = 1; count < N_PRIMES; count++) {
+ primes[count] = findNextPrimeAfter(primes[count - 1]);
+ }
+ }
+
+ /**
+ * Returns <code>sample</code>.length values chosen from the first <code>collectionSize</code>
+ * locations of <code>collection</code>, using the HASHING algorithm. Performance measurements
+ * are returned in <code>times</code>, which must be an array of at least three longs. The first
+ * will be set when the algorithm starts; the second, when a hash key has been calculated and
+ * inserted into the priority queue for every element in the collection; and the third when the
+ * original elements associated with the keys remaining in the PQ have been stored in the sample
+ * array for return.
+ * <P>
+ * This algorithm slows as the sample size becomes a significant fraction of the collection
+ * size, because the PQ is as large as the sample set, and will not do early rejection of values
+ * below the minimum until it fills up, and a larger PQ contains more small values to be purged,
+ * resulting in less early rejection and more logN insertions.
+ *
+ * @param collection The set to be sampled.
+ * @param collectionSize The number of values to use (starting from first).
+ * @param sample The array in which to return the sample.
+ * @param times The times of three events, for measuring performance.
+ */
+ private static void sample2(ScoredDocIDs collection, int collectionSize, int[] sample, long[] times)
+ throws IOException {
+ if (returnTimings) {
+ times[0] = System.currentTimeMillis();
+ }
+ int sampleSize = sample.length;
+ IntPriorityQueue pq = new IntPriorityQueue(sampleSize);
+ /*
+ * Convert every value in the collection to a hashed "weight" value, and insert
+ * into a bounded PQ (retains only sampleSize highest weights).
+ */
+ ScoredDocIDsIterator it = collection.iterator();
+ while (it.next()) {
+ pq.insertWithReuse((int)(it.getDocID() * PHI_32) & 0x7FFFFFFF);
+ }
+ if (returnTimings) {
+ times[1] = System.currentTimeMillis();
+ }
+ /*
+ * Extract heap, convert weights back to original values, and return as integers.
+ */
+ Object[] heap = pq.getHeap();
+ for (int si = 0; si < sampleSize; si++) {
+ sample[si] = (int)(((IntPriorityQueue.MI)(heap[si+1])).value * PHI_32I) & 0x7FFFFFFF;
+ }
+ if (returnTimings) {
+ times[2] = System.currentTimeMillis();
+ }
+ }
+
+ /**
+ * A bounded priority queue for Integers, to retain a specified number of
+ * the highest-weighted values for return as a random sample.
+ */
+ private static class IntPriorityQueue extends PriorityQueue<Object> {
+
+ /**
+ * Creates a bounded PQ of size <code>size</code>.
+ * @param size The number of elements to retain.
+ */
+ public IntPriorityQueue(int size) {
+ initialize(size);
+ }
+
+ /**
+ * Inserts an integer with overflow and object reuse.
+ */
+ public void insertWithReuse(int intval) {
+ if (this.mi == null) {
+ this.mi = new MI();
+ }
+ this.mi.value = intval;
+ this.mi = (MI)this.insertWithOverflow(this.mi);
+ }
+
+ /**
+ * Returns the underlying data structure for faster access. Extracting elements
+ * one at a time would require N logN time, and since we want the elements sorted
+ * in ascending order by value (not weight), the array is useful as-is.
+ * @return The underlying heap array.
+ */
+ public Object[] getHeap() {
+ return getHeapArray();
+ }
+
+ /**
+ * Returns true if <code>o1<code>'s weight is less than that of <code>o2</code>, for
+ * ordering in the PQ.
+ * @return True if <code>o1</code> weighs less than <code>o2</code>.
+ */
+ @Override
+ public boolean lessThan(Object o1, Object o2) {
+ return ((MI)o1).value < ((MI)o2).value;
+ }
+
+ /**
+ * A mutable integer that lets queue objects be reused once they start overflowing.
+ */
+ private static class MI {
+ MI() { }
+ public int value;
+ }
+
+ /**
+ * The mutable integer instance for reuse after first overflow.
+ */
+ private MI mi;
+
+ }
+
+ /**
+ * For specifying which sampling algorithm to use.
+ */
+ private enum Algorithm {
+
+ /**
+ * Specifies a methodical traversal algorithm, which is guaranteed to span the collection
+ * at least once, and never to return duplicates. Faster than the hashing algorithm and
+ * uses much less space, but the randomness of the sample may be affected by systematic
+ * variations in the collection. Requires only an array for the sample, and visits only
+ * the number of elements in the sample set, not the full set.
+ */
+ // TODO (Facet): This one produces a bimodal distribution (very flat around
+ // each peak!) for collection size 10M and sample sizes 10k and 10544.
+ // Figure out why.
+ TRAVERSAL,
+
+ /**
+ * Specifies a Fibonacci-style hash algorithm (see Knuth, S&S), which generates a less
+ * systematically distributed subset of the sampled collection than the traversal method,
+ * but requires a bounded priority queue the size of the sample, and creates an object
+ * containing a sampled value and its hash, for every element in the full set.
+ */
+ HASHING
+ }
+
+ /**
+ * For specifying whether to sort the sample.
+ */
+ private enum Sorted {
+
+ /**
+ * Sort resulting sample before returning.
+ */
+ YES,
+
+ /**
+ *Do not sort the resulting sample.
+ */
+ NO
+ }
+
+ /**
+ * Magic number 1: prime closest to phi, in 32 bits.
+ */
+ private static final long PHI_32 = 2654435769L;
+
+ /**
+ * Magic number 2: multiplicative inverse of PHI_32, modulo 2**32.
+ */
+ private static final long PHI_32I = 340573321L;
+
+ /**
+ * Switch to cause methods to return timings.
+ */
+ private static boolean returnTimings = false;
+
+}
fnp / pylucene.git / blobdiff