package org.apache.lucene.facet.search;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map.Entry;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.apache.lucene.facet.complements.TotalFacetCounts;
import org.apache.lucene.facet.complements.TotalFacetCountsCache;
import org.apache.lucene.facet.params.FacetIndexingParams;
import org.apache.lucene.facet.params.FacetSearchParams;
import org.apache.lucene.facet.partitions.IntermediateFacetResult;
import org.apache.lucene.facet.partitions.PartitionsFacetResultsHandler;
import org.apache.lucene.facet.search.FacetRequest.ResultMode;
import org.apache.lucene.facet.search.FacetsCollector.MatchingDocs;
import org.apache.lucene.facet.taxonomy.TaxonomyReader;
import org.apache.lucene.facet.util.PartitionsUtils;
import org.apache.lucene.facet.util.ScoredDocIdsUtils;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.util.IntsRef;
/*
* 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.
*/
/**
* Standard implementation for {@link FacetsAccumulator}, utilizing partitions to save on memory.
* <p>
* Why partitions? Because if there are say 100M categories out of which
* only top K are required, we must first compute value for all 100M categories
* (going over all documents) and only then could we select top K.
* This is made easier on memory by working in partitions of distinct categories:
* Once a values for a partition are found, we take the top K for that
* partition and work on the next partition, them merge the top K of both,
* and so forth, thereby computing top K with RAM needs for the size of
* a single partition rather than for the size of all the 100M categories.
* <p>
* Decision on partitions size is done at indexing time, and the facet information
* for each partition is maintained separately.
* <p>
* <u>Implementation detail:</u> Since facets information of each partition is
* maintained in a separate "category list", we can be more efficient
* at search time, because only the facet info for a single partition
* need to be read while processing that partition.
*
* @lucene.experimental
*/
public class StandardFacetsAccumulator extends FacetsAccumulator {
private static final Logger logger = Logger.getLogger(StandardFacetsAccumulator.class.getName());
/**
* Default threshold for using the complements optimization.
* If accumulating facets for a document set larger than this ratio of the index size than
* perform the complement optimization.
* @see #setComplementThreshold(double) for more info on the complements optimization.
*/
public static final double DEFAULT_COMPLEMENT_THRESHOLD = 0.6;
/**
* Passing this to {@link #setComplementThreshold(double)} will disable using complement optimization.
*/
public static final double DISABLE_COMPLEMENT = Double.POSITIVE_INFINITY; // > 1 actually
/**
* Passing this to {@link #setComplementThreshold(double)} will force using complement optimization.
*/
public static final double FORCE_COMPLEMENT = 0; // <=0
protected int partitionSize;
protected int maxPartitions;
protected boolean isUsingComplements;
private TotalFacetCounts totalFacetCounts;
private Object accumulateGuard;
private double complementThreshold = DEFAULT_COMPLEMENT_THRESHOLD;
public StandardFacetsAccumulator(FacetSearchParams searchParams, IndexReader indexReader,
TaxonomyReader taxonomyReader) {
this(searchParams, indexReader, taxonomyReader, new FacetArrays(
PartitionsUtils.partitionSize(searchParams.indexingParams, taxonomyReader)));
}
public StandardFacetsAccumulator(FacetSearchParams searchParams, IndexReader indexReader,
TaxonomyReader taxonomyReader, FacetArrays facetArrays) {
super(searchParams, indexReader, taxonomyReader, facetArrays);
// can only be computed later when docids size is known
isUsingComplements = false;
partitionSize = PartitionsUtils.partitionSize(searchParams.indexingParams, taxonomyReader);
maxPartitions = (int) Math.ceil(this.taxonomyReader.getSize() / (double) partitionSize);
accumulateGuard = new Object();
}
// TODO: this should be removed once we clean the API
public List<FacetResult> accumulate(ScoredDocIDs docids) throws IOException {
// synchronize to prevent calling two accumulate()'s at the same time.
// We decided not to synchronize the method because that might mislead
// users to feel encouraged to call this method simultaneously.
synchronized (accumulateGuard) {
// only now we can compute this
isUsingComplements = shouldComplement(docids);
if (isUsingComplements) {
try {
totalFacetCounts = TotalFacetCountsCache.getSingleton().getTotalCounts(indexReader, taxonomyReader,
searchParams.indexingParams);
if (totalFacetCounts != null) {
docids = ScoredDocIdsUtils.getComplementSet(docids, indexReader);
} else {
isUsingComplements = false;
}
} catch (UnsupportedOperationException e) {
// TODO (Facet): this exception is thrown from TotalCountsKey if the
// IndexReader used does not support getVersion(). We should re-think
// this: is this tiny detail worth disabling total counts completely
// for such readers? Currently, it's not supported by Parallel and
// MultiReader, which might be problematic for several applications.
// We could, for example, base our "isCurrent" logic on something else
// than the reader's version. Need to think more deeply about it.
if (logger.isLoggable(Level.FINEST)) {
logger.log(Level.FINEST, "IndexReader used does not support completents: ", e);
}
isUsingComplements = false;
} catch (IOException e) {
if (logger.isLoggable(Level.FINEST)) {
logger.log(Level.FINEST, "Failed to load/calculate total counts (complement counting disabled): ", e);
}
// silently fail if for some reason failed to load/save from/to dir
isUsingComplements = false;
} catch (Exception e) {
// give up: this should not happen!
throw new IOException("PANIC: Got unexpected exception while trying to get/calculate total counts", e);
}
}
docids = actualDocsToAccumulate(docids);
HashMap<FacetRequest, IntermediateFacetResult> fr2tmpRes = new HashMap<FacetRequest, IntermediateFacetResult>();
try {
for (int part = 0; part < maxPartitions; part++) {
// fill arrays from category lists
fillArraysForPartition(docids, facetArrays, part);
int offset = part * partitionSize;
// for each partition we go over all requests and handle
// each, where the request maintains the merged result.
// In this implementation merges happen after each partition,
// but other impl could merge only at the end.
final HashSet<FacetRequest> handledRequests = new HashSet<FacetRequest>();
for (FacetRequest fr : searchParams.facetRequests) {
// Handle and merge only facet requests which were not already handled.
if (handledRequests.add(fr)) {
PartitionsFacetResultsHandler frHndlr = createFacetResultsHandler(fr);
IntermediateFacetResult res4fr = frHndlr.fetchPartitionResult(offset);
IntermediateFacetResult oldRes = fr2tmpRes.get(fr);
if (oldRes != null) {
res4fr = frHndlr.mergeResults(oldRes, res4fr);
}
fr2tmpRes.put(fr, res4fr);
}
}
}
} finally {
facetArrays.free();
}
// gather results from all requests into a list for returning them
List<FacetResult> res = new ArrayList<FacetResult>();
for (FacetRequest fr : searchParams.facetRequests) {
PartitionsFacetResultsHandler frHndlr = createFacetResultsHandler(fr);
IntermediateFacetResult tmpResult = fr2tmpRes.get(fr);
if (tmpResult == null) {
// Add empty FacetResult:
res.add(emptyResult(taxonomyReader.getOrdinal(fr.categoryPath), fr));
continue;
}
FacetResult facetRes = frHndlr.renderFacetResult(tmpResult);
// final labeling if allowed (because labeling is a costly operation)
frHndlr.labelResult(facetRes);
res.add(facetRes);
}
return res;
}
}
/** check if all requests are complementable */
protected boolean mayComplement() {
for (FacetRequest freq : searchParams.facetRequests) {
if (!(freq instanceof CountFacetRequest)) {
return false;
}
}
return true;
}
@Override
protected PartitionsFacetResultsHandler createFacetResultsHandler(FacetRequest fr) {
if (fr.getResultMode() == ResultMode.PER_NODE_IN_TREE) {
return new TopKInEachNodeHandler(taxonomyReader, fr, facetArrays);
} else {
return new TopKFacetResultsHandler(taxonomyReader, fr, facetArrays);
}
}
/**
* Set the actual set of documents over which accumulation should take place.
* <p>
* Allows to override the set of documents to accumulate for. Invoked just
* before actual accumulating starts. From this point that set of documents
* remains unmodified. Default implementation just returns the input
* unchanged.
*
* @param docids
* candidate documents to accumulate for
* @return actual documents to accumulate for
*/
protected ScoredDocIDs actualDocsToAccumulate(ScoredDocIDs docids) throws IOException {
return docids;
}
/** Check if it is worth to use complements */
protected boolean shouldComplement(ScoredDocIDs docids) {
return mayComplement() && (docids.size() > indexReader.numDocs() * getComplementThreshold()) ;
}
/**
* Iterate over the documents for this partition and fill the facet arrays with the correct
* count/complement count/value.
*/
private final void fillArraysForPartition(ScoredDocIDs docids, FacetArrays facetArrays, int partition)
throws IOException {
if (isUsingComplements) {
initArraysByTotalCounts(facetArrays, partition, docids.size());
} else {
facetArrays.free(); // to get a cleared array for this partition
}
HashMap<CategoryListIterator, Aggregator> categoryLists = getCategoryListMap(facetArrays, partition);
IntsRef ordinals = new IntsRef(32); // a reasonable start capacity for most common apps
for (Entry<CategoryListIterator, Aggregator> entry : categoryLists.entrySet()) {
final ScoredDocIDsIterator iterator = docids.iterator();
final CategoryListIterator categoryListIter = entry.getKey();
final Aggregator aggregator = entry.getValue();
Iterator<AtomicReaderContext> contexts = indexReader.leaves().iterator();
AtomicReaderContext current = null;
int maxDoc = -1;
while (iterator.next()) {
int docID = iterator.getDocID();
if (docID >= maxDoc) {
boolean iteratorDone = false;
do { // find the segment which contains this document
if (!contexts.hasNext()) {
throw new RuntimeException("ScoredDocIDs contains documents outside this reader's segments !?");
}
current = contexts.next();
maxDoc = current.docBase + current.reader().maxDoc();
if (docID < maxDoc) { // segment has docs, check if it has categories
boolean validSegment = categoryListIter.setNextReader(current);
validSegment &= aggregator.setNextReader(current);
if (!validSegment) { // if categoryList or aggregtor say it's an invalid segment, skip all docs
while (docID < maxDoc && iterator.next()) {
docID = iterator.getDocID();
}
if (docID < maxDoc) {
iteratorDone = true;
}
}
}
} while (docID >= maxDoc);
if (iteratorDone) { // iterator finished, terminate the loop
break;
}
}
docID -= current.docBase;
categoryListIter.getOrdinals(docID, ordinals);
if (ordinals.length == 0) {
continue; // document does not have category ordinals
}
aggregator.aggregate(docID, iterator.getScore(), ordinals);
}
}
}
/** Init arrays for partition by total counts, optionally applying a factor */
private final void initArraysByTotalCounts(FacetArrays facetArrays, int partition, int nAccumulatedDocs) {
int[] intArray = facetArrays.getIntArray();
totalFacetCounts.fillTotalCountsForPartition(intArray, partition);
double totalCountsFactor = getTotalCountsFactor();
// fix total counts, but only if the effect of this would be meaningful.
if (totalCountsFactor < 0.99999) {
int delta = nAccumulatedDocs + 1;
for (int i = 0; i < intArray.length; i++) {
intArray[i] *= totalCountsFactor;
// also translate to prevent loss of non-positive values
// due to complement sampling (ie if sampled docs all decremented a certain category).
intArray[i] += delta;
}
}
}
/**
* Expert: factor by which counts should be multiplied when initializing
* the count arrays from total counts.
* Default implementation for this returns 1, which is a no op.
* @return a factor by which total counts should be multiplied
*/
protected double getTotalCountsFactor() {
return 1;
}
/**
* Create an {@link Aggregator} and a {@link CategoryListIterator} for each
* and every {@link FacetRequest}. Generating a map, matching each
* categoryListIterator to its matching aggregator.
* <p>
* If two CategoryListIterators are served by the same aggregator, a single
* aggregator is returned for both.
*
* <b>NOTE: </b>If a given category list iterator is needed with two different
* aggregators (e.g counting and association) - an exception is thrown as this
* functionality is not supported at this time.
*/
protected HashMap<CategoryListIterator, Aggregator> getCategoryListMap(FacetArrays facetArrays,
int partition) throws IOException {
HashMap<CategoryListIterator, Aggregator> categoryLists = new HashMap<CategoryListIterator, Aggregator>();
FacetIndexingParams indexingParams = searchParams.indexingParams;
for (FacetRequest facetRequest : searchParams.facetRequests) {
Aggregator categoryAggregator = facetRequest.createAggregator(isUsingComplements, facetArrays, taxonomyReader);
CategoryListIterator cli = indexingParams.getCategoryListParams(facetRequest.categoryPath).createCategoryListIterator(partition);
// get the aggregator
Aggregator old = categoryLists.put(cli, categoryAggregator);
if (old != null && !old.equals(categoryAggregator)) {
throw new RuntimeException("Overriding existing category list with different aggregator");
}
// if the aggregator is the same we're covered
}
return categoryLists;
}
@Override
public List<FacetResult> accumulate(List<MatchingDocs> matchingDocs) throws IOException {
return accumulate(new MatchingDocsAsScoredDocIDs(matchingDocs));
}
/**
* Returns the complement threshold.
* @see #setComplementThreshold(double)
*/
public double getComplementThreshold() {
return complementThreshold;
}
/**
* Set the complement threshold.
* This threshold will dictate whether the complements optimization is applied.
* The optimization is to count for less documents. It is useful when the same
* FacetSearchParams are used for varying sets of documents. The first time
* complements is used the "total counts" are computed - counting for all the
* documents in the collection. Then, only the complementing set of documents
* is considered, and used to decrement from the overall counts, thereby
* walking through less documents, which is faster.
* <p>
* For the default settings see {@link #DEFAULT_COMPLEMENT_THRESHOLD}.
* <p>
* To forcing complements in all cases pass {@link #FORCE_COMPLEMENT}.
* This is mostly useful for testing purposes, as forcing complements when only
* tiny fraction of available documents match the query does not make sense and
* would incur performance degradations.
* <p>
* To disable complements pass {@link #DISABLE_COMPLEMENT}.
* @param complementThreshold the complement threshold to set
* @see #getComplementThreshold()
*/
public void setComplementThreshold(double complementThreshold) {
this.complementThreshold = complementThreshold;
}
/** Returns true if complements are enabled. */
public boolean isUsingComplements() {
return isUsingComplements;
}
}