lucene-java-3.4.0/lucene/contrib/facet/src/java/org/apache/lucene/facet/search/TotalFacetCounts.java

   1 package org.apache.lucene.facet.search;
   2
   3 import java.io.BufferedInputStream;
   4 import java.io.BufferedOutputStream;
   5 import java.io.DataInputStream;
   6 import java.io.DataOutputStream;
   7 import java.io.File;
   8 import java.io.FileInputStream;
   9 import java.io.FileOutputStream;
  10 import java.io.IOException;
  11 import java.util.HashMap;
  12 import java.util.concurrent.atomic.AtomicInteger;
  13
  14 import org.apache.lucene.index.IndexReader;
  15 import org.apache.lucene.store.LockObtainFailedException;
  16
  17 import org.apache.lucene.facet.index.params.CategoryListParams;
  18 import org.apache.lucene.facet.index.params.FacetIndexingParams;
  19 import org.apache.lucene.facet.search.aggregator.Aggregator;
  20 import org.apache.lucene.facet.search.aggregator.CountingAggregator;
  21 import org.apache.lucene.facet.search.cache.CategoryListCache;
  22 import org.apache.lucene.facet.search.cache.CategoryListData;
  23 import org.apache.lucene.facet.search.params.FacetSearchParams;
  24 import org.apache.lucene.facet.taxonomy.TaxonomyReader;
  25 import org.apache.lucene.facet.util.PartitionsUtils;
  26 import org.apache.lucene.facet.util.ScoredDocIdsUtils;
  27
  28 /**
  29  * Licensed to the Apache Software Foundation (ASF) under one or more
  30  * contributor license agreements.  See the NOTICE file distributed with
  31  * this work for additional information regarding copyright ownership.
  32  * The ASF licenses this file to You under the Apache License, Version 2.0
  33  * (the "License"); you may not use this file except in compliance with
  34  * the License.  You may obtain a copy of the License at
  35  *
  36  *     http://www.apache.org/licenses/LICENSE-2.0
  37  *
  38  * Unless required by applicable law or agreed to in writing, software
  39  * distributed under the License is distributed on an "AS IS" BASIS,
  40  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  41  * See the License for the specific language governing permissions and
  42  * limitations under the License.
  43  */
  44
  45 /**
  46  * Maintain Total Facet Counts per partition, for given parameters:
  47  * <ul>
  48  *  <li>Index reader of an index</li>
  49  *  <li>Taxonomy index reader</li>
  50  *  <li>Facet indexing params (and particularly the category list params)</li>
  51  *  <li></li>
  52  * </ul>
  53  * The total facet counts are maintained as an array of arrays of integers,
  54  * where a separate array is kept for each partition.
  55  *
  56  * @lucene.experimental
  57  */
  58 public class TotalFacetCounts {
  59
  60   /** total facet counts per partition: totalCounts[partition][ordinal%partitionLength] */
  61   private int[][] totalCounts = null;
  62
  63   private final TaxonomyReader taxonomy;
  64   private final FacetIndexingParams facetIndexingParams;
  65
  66   private final static AtomicInteger atomicGen4Test = new AtomicInteger(1);
  67   /** Creation type for test purposes */
  68   enum CreationType { Computed, Loaded } // for testing
  69   final int gen4test;
  70   final CreationType createType4test;
  71
  72   /**
  73    * Construct by key - from index Directory or by recomputing.
  74    * @param key the key mapping of this total facet counts (index, taxonomy, category lists...)
  75    */
  76   private TotalFacetCounts (TaxonomyReader taxonomy, FacetIndexingParams facetIndexingParams,
  77       int[][] counts, CreationType createType4Test) throws IOException, LockObtainFailedException {
  78     this.taxonomy = taxonomy;
  79     this.facetIndexingParams = facetIndexingParams;
  80     this.totalCounts = counts;
  81     this.createType4test = createType4Test;
  82     this.gen4test = atomicGen4Test.incrementAndGet();
  83   }
  84
  85   /**
  86    * Fill a partition's array with the TotalCountsArray values.
  87    * @param partitionArray array to fill
  88    * @param partition number of required partition
  89    */
  90   public void fillTotalCountsForPartition(int[] partitionArray, int partition) {
  91     int partitionSize = partitionArray.length;
  92     int[] countArray = totalCounts[partition];
  93     if (countArray == null) {
  94       countArray = new int[partitionSize];
  95       totalCounts[partition] = countArray;
  96     }
  97     int length = Math.min(partitionSize, countArray.length);
  98     System.arraycopy(countArray, 0, partitionArray, 0, length);
  99   }
 100
 101   /**
 102    * Return the total count of an input category
 103    * @param ordinal ordinal of category whose total count is required
 104    */
 105   public int getTotalCount(int ordinal) {
 106     int partition = PartitionsUtils.partitionNumber(facetIndexingParams,ordinal);
 107     int offset = ordinal % PartitionsUtils.partitionSize(facetIndexingParams, taxonomy);
 108     return totalCounts[partition][offset];
 109   }
 110
 111   static TotalFacetCounts loadFromFile(File inputFile, TaxonomyReader taxonomy,
 112       FacetIndexingParams facetIndexingParams) throws IOException {
 113     DataInputStream dis = new DataInputStream(new BufferedInputStream(new FileInputStream(inputFile)));
 114     try {
 115       int[][] counts = new int[dis.readInt()][];
 116       for (int i=0; i<counts.length; i++) {
 117         int size = dis.readInt();
 118         if (size<0) {
 119           counts[i] = null;
 120         } else {
 121           counts[i] = new int[size];
 122           for (int j=0; j<size; j++) {
 123             counts[i][j] = dis.readInt();
 124           }
 125         }
 126       }
 127       return new TotalFacetCounts(taxonomy, facetIndexingParams, counts, CreationType.Loaded);
 128     } finally {
 129       dis.close();
 130     }
 131   }
 132
 133   static void storeToFile(File outputFile, TotalFacetCounts tfc) throws IOException {
 134     DataOutputStream dos = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(outputFile)));
 135     try {
 136       dos.writeInt(tfc.totalCounts.length);
 137       for (int[] counts : tfc.totalCounts) {
 138         if (counts == null) {
 139           dos.writeInt(-1);
 140         } else {
 141           dos.writeInt(counts.length);
 142           for (int i : counts) {
 143             dos.writeInt(i);
 144           }
 145         }
 146       }
 147     } finally {
 148       dos.close();
 149     }
 150   }
 151
 152   static TotalFacetCounts compute(final IndexReader indexReader,
 153       final TaxonomyReader taxonomy, final FacetIndexingParams facetIndexingParams,
 154       final CategoryListCache clCache) throws IOException {
 155     int partitionSize = PartitionsUtils.partitionSize(facetIndexingParams, taxonomy);
 156     final int[][] counts = new int[(int) Math.ceil(taxonomy.getSize()  /(float) partitionSize)][partitionSize];
 157     FacetSearchParams newSearchParams = new FacetSearchParams(facetIndexingParams);
 158       //createAllListsSearchParams(facetIndexingParams,  this.totalCounts);
 159     FacetsAccumulator fe = new StandardFacetsAccumulator(newSearchParams, indexReader, taxonomy) {
 160       @Override
 161       protected HashMap<CategoryListIterator, Aggregator> getCategoryListMap(
 162           FacetArrays facetArrays, int partition) throws IOException {
 163
 164         Aggregator aggregator = new CountingAggregator(counts[partition]);
 165         HashMap<CategoryListIterator, Aggregator> map = new HashMap<CategoryListIterator, Aggregator>();
 166         for (CategoryListParams clp: facetIndexingParams.getAllCategoryListParams()) {
 167           final CategoryListIterator cli = clIteraor(clCache, clp, indexReader, partition);
 168           map.put(cli, aggregator);
 169         }
 170         return map;
 171       }
 172     };
 173     fe.setComplementThreshold(FacetsAccumulator.DISABLE_COMPLEMENT);
 174     fe.accumulate(ScoredDocIdsUtils.createAllDocsScoredDocIDs(indexReader));
 175     return new TotalFacetCounts(taxonomy, facetIndexingParams, counts, CreationType.Computed);
 176   }
 177
 178   static CategoryListIterator clIteraor(CategoryListCache clCache, CategoryListParams clp,
 179       IndexReader indexReader, int partition) throws IOException {
 180     if (clCache != null) {
 181       CategoryListData cld = clCache.get(clp);
 182       if (cld != null) {
 183         return cld.iterator(partition);
 184       }
 185     }
 186     return clp.createCategoryListIterator(indexReader, partition);
 187   }
 188 }