+++ /dev/null
-package org.apache.lucene.search;
-
-/**
- * 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.
- */
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-
-import org.apache.lucene.index.IndexReader;
-import org.apache.lucene.util.RamUsageEstimator;
-
-/**
- * Caches all docs, and optionally also scores, coming from
- * a search, and is then able to replay them to another
- * collector. You specify the max RAM this class may use.
- * Once the collection is done, call {@link #isCached}. If
- * this returns true, you can use {@link #replay} against a
- * new collector. If it returns false, this means too much
- * RAM was required and you must instead re-run the original
- * search.
- *
- * <p><b>NOTE</b>: this class consumes 4 (or 8 bytes, if
- * scoring is cached) per collected document. If the result
- * set is large this can easily be a very substantial amount
- * of RAM!
- *
- * <p><b>NOTE</b>: this class caches at least 128 documents
- * before checking RAM limits.
- *
- * <p>See the Lucene <tt>contrib/grouping</tt> module for more
- * details including a full code example.</p>
- *
- * @lucene.experimental
- */
-public abstract class CachingCollector extends Collector {
-
- // Max out at 512K arrays
- private static final int MAX_ARRAY_SIZE = 512 * 1024;
- private static final int INITIAL_ARRAY_SIZE = 128;
- private final static int[] EMPTY_INT_ARRAY = new int[0];
-
- private static class SegStart {
- public final IndexReader reader;
- public final int base;
- public final int end;
-
- public SegStart(IndexReader reader, int base, int end) {
- this.reader = reader;
- this.base = base;
- this.end = end;
- }
- }
-
- private static final class CachedScorer extends Scorer {
-
- // NOTE: these members are package-private b/c that way accessing them from
- // the outer class does not incur access check by the JVM. The same
- // situation would be if they were defined in the outer class as private
- // members.
- int doc;
- float score;
-
- private CachedScorer() { super((Weight) null); }
-
- @Override
- public final float score() { return score; }
-
- @Override
- public final int advance(int target) { throw new UnsupportedOperationException(); }
-
- @Override
- public final int docID() { return doc; }
-
- @Override
- public final float freq() { throw new UnsupportedOperationException(); }
-
- @Override
- public final int nextDoc() { throw new UnsupportedOperationException(); }
- }
-
- // A CachingCollector which caches scores
- private static final class ScoreCachingCollector extends CachingCollector {
-
- private final CachedScorer cachedScorer;
- private final List<float[]> cachedScores;
-
- private Scorer scorer;
- private float[] curScores;
-
- ScoreCachingCollector(Collector other, double maxRAMMB) {
- super(other, maxRAMMB, true);
-
- cachedScorer = new CachedScorer();
- cachedScores = new ArrayList<float[]>();
- curScores = new float[128];
- cachedScores.add(curScores);
- }
-
- ScoreCachingCollector(Collector other, int maxDocsToCache) {
- super(other, maxDocsToCache);
-
- cachedScorer = new CachedScorer();
- cachedScores = new ArrayList<float[]>();
- curScores = new float[INITIAL_ARRAY_SIZE];
- cachedScores.add(curScores);
- }
-
- @Override
- public void collect(int doc) throws IOException {
-
- if (curDocs == null) {
- // Cache was too large
- cachedScorer.score = scorer.score();
- cachedScorer.doc = doc;
- other.collect(doc);
- return;
- }
-
- // Allocate a bigger array or abort caching
- if (upto == curDocs.length) {
- base += upto;
-
- // Compute next array length - don't allocate too big arrays
- int nextLength = 8*curDocs.length;
- if (nextLength > MAX_ARRAY_SIZE) {
- nextLength = MAX_ARRAY_SIZE;
- }
-
- if (base + nextLength > maxDocsToCache) {
- // try to allocate a smaller array
- nextLength = maxDocsToCache - base;
- if (nextLength <= 0) {
- // Too many docs to collect -- clear cache
- curDocs = null;
- curScores = null;
- cachedSegs.clear();
- cachedDocs.clear();
- cachedScores.clear();
- cachedScorer.score = scorer.score();
- cachedScorer.doc = doc;
- other.collect(doc);
- return;
- }
- }
-
- curDocs = new int[nextLength];
- cachedDocs.add(curDocs);
- curScores = new float[nextLength];
- cachedScores.add(curScores);
- upto = 0;
- }
-
- curDocs[upto] = doc;
- cachedScorer.score = curScores[upto] = scorer.score();
- upto++;
- cachedScorer.doc = doc;
- other.collect(doc);
- }
-
- @Override
- public void replay(Collector other) throws IOException {
- replayInit(other);
-
- int curUpto = 0;
- int curBase = 0;
- int chunkUpto = 0;
- curDocs = EMPTY_INT_ARRAY;
- for (SegStart seg : cachedSegs) {
- other.setNextReader(seg.reader, seg.base);
- other.setScorer(cachedScorer);
- while (curBase + curUpto < seg.end) {
- if (curUpto == curDocs.length) {
- curBase += curDocs.length;
- curDocs = cachedDocs.get(chunkUpto);
- curScores = cachedScores.get(chunkUpto);
- chunkUpto++;
- curUpto = 0;
- }
- cachedScorer.score = curScores[curUpto];
- cachedScorer.doc = curDocs[curUpto];
- other.collect(curDocs[curUpto++]);
- }
- }
- }
-
- @Override
- public void setScorer(Scorer scorer) throws IOException {
- this.scorer = scorer;
- other.setScorer(cachedScorer);
- }
-
- @Override
- public String toString() {
- if (isCached()) {
- return "CachingCollector (" + (base+upto) + " docs & scores cached)";
- } else {
- return "CachingCollector (cache was cleared)";
- }
- }
-
- }
-
- // A CachingCollector which does not cache scores
- private static final class NoScoreCachingCollector extends CachingCollector {
-
- NoScoreCachingCollector(Collector other, double maxRAMMB) {
- super(other, maxRAMMB, false);
- }
-
- NoScoreCachingCollector(Collector other, int maxDocsToCache) {
- super(other, maxDocsToCache);
- }
-
- @Override
- public void collect(int doc) throws IOException {
-
- if (curDocs == null) {
- // Cache was too large
- other.collect(doc);
- return;
- }
-
- // Allocate a bigger array or abort caching
- if (upto == curDocs.length) {
- base += upto;
-
- // Compute next array length - don't allocate too big arrays
- int nextLength = 8*curDocs.length;
- if (nextLength > MAX_ARRAY_SIZE) {
- nextLength = MAX_ARRAY_SIZE;
- }
-
- if (base + nextLength > maxDocsToCache) {
- // try to allocate a smaller array
- nextLength = maxDocsToCache - base;
- if (nextLength <= 0) {
- // Too many docs to collect -- clear cache
- curDocs = null;
- cachedSegs.clear();
- cachedDocs.clear();
- other.collect(doc);
- return;
- }
- }
-
- curDocs = new int[nextLength];
- cachedDocs.add(curDocs);
- upto = 0;
- }
-
- curDocs[upto] = doc;
- upto++;
- other.collect(doc);
- }
-
- @Override
- public void replay(Collector other) throws IOException {
- replayInit(other);
-
- int curUpto = 0;
- int curbase = 0;
- int chunkUpto = 0;
- curDocs = EMPTY_INT_ARRAY;
- for (SegStart seg : cachedSegs) {
- other.setNextReader(seg.reader, seg.base);
- while (curbase + curUpto < seg.end) {
- if (curUpto == curDocs.length) {
- curbase += curDocs.length;
- curDocs = cachedDocs.get(chunkUpto);
- chunkUpto++;
- curUpto = 0;
- }
- other.collect(curDocs[curUpto++]);
- }
- }
- }
-
- @Override
- public void setScorer(Scorer scorer) throws IOException {
- other.setScorer(scorer);
- }
-
- @Override
- public String toString() {
- if (isCached()) {
- return "CachingCollector (" + (base+upto) + " docs cached)";
- } else {
- return "CachingCollector (cache was cleared)";
- }
- }
-
- }
-
- // TODO: would be nice if a collector defined a
- // needsScores() method so we can specialize / do checks
- // up front. This is only relevant for the ScoreCaching
- // version -- if the wrapped Collector does not need
- // scores, it can avoid cachedScorer entirely.
- protected final Collector other;
-
- protected final int maxDocsToCache;
- protected final List<SegStart> cachedSegs = new ArrayList<SegStart>();
- protected final List<int[]> cachedDocs;
-
- private IndexReader lastReader;
-
- protected int[] curDocs;
- protected int upto;
- protected int base;
- protected int lastDocBase;
-
- /**
- * Creates a {@link CachingCollector} which does not wrap another collector.
- * The cached documents and scores can later be {@link #replay(Collector)
- * replayed}.
- *
- * @param acceptDocsOutOfOrder
- * whether documents are allowed to be collected out-of-order
- */
- public static CachingCollector create(final boolean acceptDocsOutOfOrder, boolean cacheScores, double maxRAMMB) {
- Collector other = new Collector() {
- @Override
- public boolean acceptsDocsOutOfOrder() {
- return acceptDocsOutOfOrder;
- }
-
- @Override
- public void setScorer(Scorer scorer) throws IOException {}
-
- @Override
- public void collect(int doc) throws IOException {}
-
- @Override
- public void setNextReader(IndexReader reader, int docBase) throws IOException {}
-
- };
- return create(other, cacheScores, maxRAMMB);
- }
-
- /**
- * Create a new {@link CachingCollector} that wraps the given collector and
- * caches documents and scores up to the specified RAM threshold.
- *
- * @param other
- * the Collector to wrap and delegate calls to.
- * @param cacheScores
- * whether to cache scores in addition to document IDs. Note that
- * this increases the RAM consumed per doc
- * @param maxRAMMB
- * the maximum RAM in MB to consume for caching the documents and
- * scores. If the collector exceeds the threshold, no documents and
- * scores are cached.
- */
- public static CachingCollector create(Collector other, boolean cacheScores, double maxRAMMB) {
- return cacheScores ? new ScoreCachingCollector(other, maxRAMMB) : new NoScoreCachingCollector(other, maxRAMMB);
- }
-
- /**
- * Create a new {@link CachingCollector} that wraps the given collector and
- * caches documents and scores up to the specified max docs threshold.
- *
- * @param other
- * the Collector to wrap and delegate calls to.
- * @param cacheScores
- * whether to cache scores in addition to document IDs. Note that
- * this increases the RAM consumed per doc
- * @param maxDocsToCache
- * the maximum number of documents for caching the documents and
- * possible the scores. If the collector exceeds the threshold,
- * no documents and scores are cached.
- */
- public static CachingCollector create(Collector other, boolean cacheScores, int maxDocsToCache) {
- return cacheScores ? new ScoreCachingCollector(other, maxDocsToCache) : new NoScoreCachingCollector(other, maxDocsToCache);
- }
-
- // Prevent extension from non-internal classes
- private CachingCollector(Collector other, double maxRAMMB, boolean cacheScores) {
- this.other = other;
-
- cachedDocs = new ArrayList<int[]>();
- curDocs = new int[INITIAL_ARRAY_SIZE];
- cachedDocs.add(curDocs);
-
- int bytesPerDoc = RamUsageEstimator.NUM_BYTES_INT;
- if (cacheScores) {
- bytesPerDoc += RamUsageEstimator.NUM_BYTES_FLOAT;
- }
- maxDocsToCache = (int) ((maxRAMMB * 1024 * 1024) / bytesPerDoc);
- }
-
- private CachingCollector(Collector other, int maxDocsToCache) {
- this.other = other;
-
- cachedDocs = new ArrayList<int[]>();
- curDocs = new int[INITIAL_ARRAY_SIZE];
- cachedDocs.add(curDocs);
- this.maxDocsToCache = maxDocsToCache;
- }
-
- @Override
- public boolean acceptsDocsOutOfOrder() {
- return other.acceptsDocsOutOfOrder();
- }
-
- public boolean isCached() {
- return curDocs != null;
- }
-
- @Override
- public void setNextReader(IndexReader reader, int docBase) throws IOException {
- other.setNextReader(reader, docBase);
- if (lastReader != null) {
- cachedSegs.add(new SegStart(lastReader, lastDocBase, base + upto));
- }
- lastDocBase = docBase;
- lastReader = reader;
- }
-
- /** Reused by the specialized inner classes. */
- void replayInit(Collector other) {
- if (!isCached()) {
- throw new IllegalStateException("cannot replay: cache was cleared because too much RAM was required");
- }
-
- if (!other.acceptsDocsOutOfOrder() && this.other.acceptsDocsOutOfOrder()) {
- throw new IllegalArgumentException(
- "cannot replay: given collector does not support "
- + "out-of-order collection, while the wrapped collector does. "
- + "Therefore cached documents may be out-of-order.");
- }
-
- //System.out.println("CC: replay totHits=" + (upto + base));
- if (lastReader != null) {
- cachedSegs.add(new SegStart(lastReader, lastDocBase, base+upto));
- lastReader = null;
- }
- }
-
- /**
- * Replays the cached doc IDs (and scores) to the given Collector. If this
- * instance does not cache scores, then Scorer is not set on
- * {@code other.setScorer} as well as scores are not replayed.
- *
- * @throws IllegalStateException
- * if this collector is not cached (i.e., if the RAM limits were too
- * low for the number of documents + scores to cache).
- * @throws IllegalArgumentException
- * if the given Collect's does not support out-of-order collection,
- * while the collector passed to the ctor does.
- */
- public abstract void replay(Collector other) throws IOException;
-
-}
fnp / pylucene.git / blobdiff