lucene-java-3.5.0/lucene/contrib/facet/src/test/org/apache/lucene/facet/FacetTestUtils.java

   1 package org.apache.lucene.facet;
   2
   3 import java.io.IOException;
   4 import java.util.ArrayList;
   5 import java.util.Collection;
   6
   7 import org.apache.lucene.analysis.standard.StandardAnalyzer;
   8 import org.apache.lucene.document.Document;
   9 import org.apache.lucene.document.Field;
  10 import org.apache.lucene.document.Field.Index;
  11 import org.apache.lucene.document.Field.Store;
  12 import org.apache.lucene.document.Field.TermVector;
  13 import org.apache.lucene.index.CorruptIndexException;
  14 import org.apache.lucene.index.IndexReader;
  15 import org.apache.lucene.index.IndexWriter;
  16 import org.apache.lucene.index.IndexWriterConfig;
  17 import org.apache.lucene.index.RandomIndexWriter;
  18 import org.apache.lucene.search.Collector;
  19 import org.apache.lucene.search.IndexSearcher;
  20 import org.apache.lucene.search.MatchAllDocsQuery;
  21 import org.apache.lucene.search.TopScoreDocCollector;
  22 import org.apache.lucene.store.Directory;
  23
  24 import org.apache.lucene.search.MultiCollector;
  25 import org.apache.lucene.util.LuceneTestCase;
  26 import org.apache.lucene.facet.index.CategoryDocumentBuilder;
  27 import org.apache.lucene.facet.index.params.DefaultFacetIndexingParams;
  28 import org.apache.lucene.facet.index.params.FacetIndexingParams;
  29 import org.apache.lucene.facet.search.FacetsCollector;
  30 import org.apache.lucene.facet.search.params.CountFacetRequest;
  31 import org.apache.lucene.facet.search.params.FacetRequest;
  32 import org.apache.lucene.facet.search.params.FacetSearchParams;
  33 import org.apache.lucene.facet.taxonomy.CategoryPath;
  34 import org.apache.lucene.facet.taxonomy.TaxonomyReader;
  35 import org.apache.lucene.facet.taxonomy.TaxonomyWriter;
  36 import org.apache.lucene.facet.taxonomy.directory.DirectoryTaxonomyReader;
  37 import org.apache.lucene.facet.taxonomy.directory.DirectoryTaxonomyWriter;
  38
  39 /**
  40  * Licensed to the Apache Software Foundation (ASF) under one or more
  41  * contributor license agreements.  See the NOTICE file distributed with
  42  * this work for additional information regarding copyright ownership.
  43  * The ASF licenses this file to You under the Apache License, Version 2.0
  44  * (the "License"); you may not use this file except in compliance with
  45  * the License.  You may obtain a copy of the License at
  46  *
  47  *     http://www.apache.org/licenses/LICENSE-2.0
  48  *
  49  * Unless required by applicable law or agreed to in writing, software
  50  * distributed under the License is distributed on an "AS IS" BASIS,
  51  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  52  * See the License for the specific language governing permissions and
  53  * limitations under the License.
  54  */
  55
  56 public class FacetTestUtils {
  57
  58   public static Directory[][] createIndexTaxonomyDirs(int number) throws IOException {
  59     Directory[][] dirs = new Directory[number][2];
  60     for (int i = 0; i < number; i++) {
  61       dirs[i][0] = LuceneTestCase.newDirectory();
  62       dirs[i][1] = LuceneTestCase.newDirectory();
  63     }
  64     return dirs;
  65   }
  66
  67   public static IndexTaxonomyReaderPair[] createIndexTaxonomyReaderPair(
  68       Directory[][] dirs) throws IOException {
  69     IndexTaxonomyReaderPair[] pairs = new IndexTaxonomyReaderPair[dirs.length];
  70     for (int i = 0; i < dirs.length; i++) {
  71       IndexTaxonomyReaderPair pair = new IndexTaxonomyReaderPair();
  72       pair.indexReader = IndexReader.open(dirs[i][0]);
  73       pair.indexSearcher = new IndexSearcher(pair.indexReader);
  74       pair.taxReader = new DirectoryTaxonomyReader(dirs[i][1]);
  75       pairs[i] = pair;
  76     }
  77     return pairs;
  78   }
  79
  80   public static IndexTaxonomyWriterPair[] createIndexTaxonomyWriterPair(
  81       Directory[][] dirs) throws IOException {
  82     IndexTaxonomyWriterPair[] pairs = new IndexTaxonomyWriterPair[dirs.length];
  83     for (int i = 0; i < dirs.length; i++) {
  84       IndexTaxonomyWriterPair pair = new IndexTaxonomyWriterPair();
  85       pair.indexWriter = new IndexWriter(dirs[i][0], new IndexWriterConfig(
  86           LuceneTestCase.TEST_VERSION_CURRENT, new StandardAnalyzer(
  87               LuceneTestCase.TEST_VERSION_CURRENT)));
  88       pair.taxWriter = new DirectoryTaxonomyWriter(dirs[i][1]);
  89       pair.indexWriter.commit();
  90       pair.taxWriter.commit();
  91       pairs[i] = pair;
  92     }
  93     return pairs;
  94   }
  95
  96   public static Collector[] search(IndexSearcher searcher,
  97       TaxonomyReader taxonomyReader, DefaultFacetIndexingParams iParams,
  98       int k, String... facetNames) throws IOException,
  99       IllegalAccessException, InstantiationException {
 100
 101     Collector[] collectors = new Collector[2];
 102
 103     FacetSearchParams facetSearchParams = new FacetSearchParams(iParams);
 104     Collection<FacetRequest> fRequests = new ArrayList<FacetRequest>();
 105     for (String facetName : facetNames) {
 106       CategoryPath cp = new CategoryPath(facetName);
 107       FacetRequest fq = new CountFacetRequest(cp, k);
 108       facetSearchParams.addFacetRequest(fq);
 109       fRequests.add(fq);
 110     }
 111
 112     TopScoreDocCollector topDocsCollector = TopScoreDocCollector.create(
 113         searcher.getIndexReader().maxDoc(), true);
 114     FacetsCollector facetsCollector = new FacetsCollector(
 115         facetSearchParams, searcher.getIndexReader(), taxonomyReader);
 116     Collector mColl = MultiCollector.wrap(topDocsCollector, facetsCollector);
 117
 118     collectors[0] = topDocsCollector;
 119     collectors[1] = facetsCollector;
 120
 121     searcher.search(new MatchAllDocsQuery(), mColl);
 122     return collectors;
 123   }
 124
 125   public static void add(FacetIndexingParams iParams, RandomIndexWriter iw,
 126       TaxonomyWriter tw, String... strings) throws IOException,
 127       CorruptIndexException {
 128     ArrayList<CategoryPath> cps = new ArrayList<CategoryPath>();
 129     CategoryPath cp = new CategoryPath(strings);
 130     cps.add(cp);
 131     Document d = new Document();
 132     new CategoryDocumentBuilder(tw, iParams).setCategoryPaths(cps).build(d);
 133     d.add(new Field("content", "alpha", Store.YES, Index.ANALYZED,
 134         TermVector.NO));
 135     iw.addDocument(d);
 136   }
 137
 138   public static class IndexTaxonomyReaderPair {
 139     public IndexReader indexReader;
 140     public TaxonomyReader taxReader;
 141     public IndexSearcher indexSearcher;
 142
 143     public void close() throws IOException {
 144       indexSearcher.close();
 145       indexReader.close();
 146       taxReader.close();
 147     }
 148
 149   }
 150
 151   public static class IndexTaxonomyWriterPair {
 152     public IndexWriter indexWriter;
 153     public TaxonomyWriter taxWriter;
 154
 155     public void close() throws IOException {
 156       indexWriter.close();
 157       taxWriter.close();
 158     }
 159
 160     public void commit() throws IOException {
 161       indexWriter.commit();
 162       taxWriter.commit();
 163     }
 164   }
 165
 166 }