lucene-java-3.5.0/lucene/contrib/benchmark/src/java/org/apache/lucene/benchmark/byTask/feeds/TrecContentSource.java

   1 package org.apache.lucene.benchmark.byTask.feeds;
   2
   3 /**
   4  * Licensed to the Apache Software Foundation (ASF) under one or more
   5  * contributor license agreements.  See the NOTICE file distributed with
   6  * this work for additional information regarding copyright ownership.
   7  * The ASF licenses this file to You under the Apache License, Version 2.0
   8  * (the "License"); you may not use this file except in compliance with
   9  * the License.  You may obtain a copy of the License at
  10  *
  11  *     http://www.apache.org/licenses/LICENSE-2.0
  12  *
  13  * Unless required by applicable law or agreed to in writing, software
  14  * distributed under the License is distributed on an "AS IS" BASIS,
  15  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  16  * See the License for the specific language governing permissions and
  17  * limitations under the License.
  18  */
  19
  20 import java.io.BufferedReader;
  21 import java.io.File;
  22 import java.io.IOException;
  23 import java.io.InputStream;
  24 import java.io.InputStreamReader;
  25 import java.io.Reader;
  26 import java.text.DateFormat;
  27 import java.text.ParsePosition;
  28 import java.text.SimpleDateFormat;
  29 import java.util.ArrayList;
  30 import java.util.Date;
  31 import java.util.Locale;
  32
  33 import org.apache.lucene.benchmark.byTask.feeds.TrecDocParser.ParsePathType;
  34 import org.apache.lucene.benchmark.byTask.utils.Config;
  35 import org.apache.lucene.benchmark.byTask.utils.StreamUtils;
  36 import org.apache.lucene.benchmark.byTask.utils.StringBuilderReader;
  37 import org.apache.lucene.util.ThreadInterruptedException;
  38
  39 /**
  40  * Implements a {@link ContentSource} over the TREC collection.
  41  * <p>
  42  * Supports the following configuration parameters (on top of
  43  * {@link ContentSource}):
  44  * <ul>
  45  * <li><b>work.dir</b> - specifies the working directory. Required if "docs.dir"
  46  * denotes a relative path (<b>default=work</b>).
  47  * <li><b>docs.dir</b> - specifies the directory where the TREC files reside.
  48  * Can be set to a relative path if "work.dir" is also specified
  49  * (<b>default=trec</b>).
  50  * <li><b>trec.doc.parser</b> - specifies the {@link TrecDocParser} class to use for
  51  * parsing the TREC documents content (<b>default=TrecGov2Parser</b>).
  52  * <li><b>html.parser</b> - specifies the {@link HTMLParser} class to use for
  53  * parsing the HTML parts of the TREC documents content (<b>default=DemoHTMLParser</b>).
  54  * <li><b>content.source.encoding</b> - if not specified, ISO-8859-1 is used.
  55  * <li><b>content.source.excludeIteration</b> - if true, do not append iteration number to docname
  56  * </ul>
  57  */
  58 public class TrecContentSource extends ContentSource {
  59
  60   private static final class DateFormatInfo {
  61     DateFormatInfo() {}
  62     DateFormat[] dfs;
  63     ParsePosition pos;
  64   }
  65
  66   public static final String DOCNO = "<DOCNO>";
  67   public static final String TERMINATING_DOCNO = "</DOCNO>";
  68   public static final String DOC = "<DOC>";
  69   public static final String TERMINATING_DOC = "</DOC>";
  70
  71   /** separator between lines in the byffer */
  72   public static final String NEW_LINE = System.getProperty("line.separator");
  73
  74   private static final String DATE_FORMATS [] = {
  75        "EEE, dd MMM yyyy kk:mm:ss z",   // Tue, 09 Dec 2003 22:39:08 GMT
  76        "EEE MMM dd kk:mm:ss yyyy z",    // Tue Dec 09 16:45:08 2003 EST
  77        "EEE, dd-MMM-':'y kk:mm:ss z",   // Tue, 09 Dec 2003 22:39:08 GMT
  78        "EEE, dd-MMM-yyy kk:mm:ss z",    // Tue, 09 Dec 2003 22:39:08 GMT
  79        "EEE MMM dd kk:mm:ss yyyy",      // Tue Dec 09 16:45:08 2003
  80        "dd MMM yyyy",                   // 1 March 1994
  81        "MMM dd, yyyy",                  // February 3, 1994
  82        "yyMMdd",                        // 910513
  83        "hhmm z.z.z. MMM dd, yyyy",       // 0901 u.t.c. April 28, 1994
  84   };
  85
  86   private ThreadLocal<DateFormatInfo> dateFormats = new ThreadLocal<DateFormatInfo>();
  87   private ThreadLocal<StringBuilderReader> trecDocReader = new ThreadLocal<StringBuilderReader>();
  88   private ThreadLocal<StringBuilder> trecDocBuffer = new ThreadLocal<StringBuilder>();
  89   private File dataDir = null;
  90   private ArrayList<File> inputFiles = new ArrayList<File>();
  91   private int nextFile = 0;
  92   private int rawDocSize = 0;
  93
  94   // Use to synchronize threads on reading from the TREC documents.
  95   private Object lock = new Object();
  96
  97   // Required for test
  98   BufferedReader reader;
  99   int iteration = 0;
 100   HTMLParser htmlParser;
 101
 102   private boolean excludeDocnameIteration;
 103   private TrecDocParser trecDocParser = new TrecGov2Parser(); // default
 104   ParsePathType currPathType; // not private for tests
 105
 106   private DateFormatInfo getDateFormatInfo() {
 107     DateFormatInfo dfi = dateFormats.get();
 108     if (dfi == null) {
 109       dfi = new DateFormatInfo();
 110       dfi.dfs = new SimpleDateFormat[DATE_FORMATS.length];
 111       for (int i = 0; i < dfi.dfs.length; i++) {
 112         dfi.dfs[i] = new SimpleDateFormat(DATE_FORMATS[i], Locale.US);
 113         dfi.dfs[i].setLenient(true);
 114       }
 115       dfi.pos = new ParsePosition(0);
 116       dateFormats.set(dfi);
 117     }
 118     return dfi;
 119   }
 120
 121   private StringBuilder getDocBuffer() {
 122     StringBuilder sb = trecDocBuffer.get();
 123     if (sb == null) {
 124       sb = new StringBuilder();
 125       trecDocBuffer.set(sb);
 126     }
 127     return sb;
 128   }
 129
 130   Reader getTrecDocReader(StringBuilder docBuffer) {
 131     StringBuilderReader r = trecDocReader.get();
 132     if (r == null) {
 133       r = new StringBuilderReader(docBuffer);
 134       trecDocReader.set(r);
 135     } else {
 136       r.set(docBuffer);
 137     }
 138     return r;
 139   }
 140
 141   HTMLParser getHtmlParser() {
 142     return htmlParser;
 143   }
 144
 145   /**
 146    * Read until a line starting with the specified <code>lineStart</code>.
 147    * @param buf buffer for collecting the data if so specified/
 148    * @param lineStart line start to look for, must not be null.
 149    * @param collectMatchLine whether to collect the matching line into <code>buffer</code>.
 150    * @param collectAll whether to collect all lines into <code>buffer</code>.
 151    * @throws IOException
 152    * @throws NoMoreDataException
 153    */
 154    private void read(StringBuilder buf, String lineStart,
 155        boolean collectMatchLine, boolean collectAll) throws IOException, NoMoreDataException {
 156     String sep = "";
 157     while (true) {
 158       String line = reader.readLine();
 159
 160       if (line == null) {
 161         openNextFile();
 162         continue;
 163       }
 164
 165       rawDocSize += line.length();
 166
 167       if (lineStart!=null && line.startsWith(lineStart)) {
 168         if (collectMatchLine) {
 169           buf.append(sep).append(line);
 170           sep = NEW_LINE;
 171         }
 172         return;
 173       }
 174
 175       if (collectAll) {
 176         buf.append(sep).append(line);
 177         sep = NEW_LINE;
 178       }
 179     }
 180   }
 181
 182   void openNextFile() throws NoMoreDataException, IOException {
 183     close();
 184     currPathType = null;
 185     while (true) {
 186       if (nextFile >= inputFiles.size()) {
 187         // exhausted files, start a new round, unless forever set to false.
 188         if (!forever) {
 189           throw new NoMoreDataException();
 190         }
 191         nextFile = 0;
 192         iteration++;
 193       }
 194       File f = inputFiles.get(nextFile++);
 195       if (verbose) {
 196         System.out.println("opening: " + f + " length: " + f.length());
 197       }
 198       try {
 199         InputStream inputStream = StreamUtils.inputStream(f); // support either gzip, bzip2, or regular text file, by extension
 200         reader = new BufferedReader(new InputStreamReader(inputStream, encoding), StreamUtils.BUFFER_SIZE);
 201         currPathType = TrecDocParser.pathType(f);
 202         return;
 203       } catch (Exception e) {
 204         if (verbose) {
 205           System.out.println("Skipping 'bad' file " + f.getAbsolutePath()+" due to "+e.getMessage());
 206           continue;
 207         }
 208         throw new NoMoreDataException();
 209       }
 210     }
 211   }
 212
 213   public Date parseDate(String dateStr) {
 214     dateStr = dateStr.trim();
 215     DateFormatInfo dfi = getDateFormatInfo();
 216     for (int i = 0; i < dfi.dfs.length; i++) {
 217       DateFormat df = dfi.dfs[i];
 218       dfi.pos.setIndex(0);
 219       dfi.pos.setErrorIndex(-1);
 220       Date d = df.parse(dateStr, dfi.pos);
 221       if (d != null) {
 222         // Parse succeeded.
 223         return d;
 224       }
 225     }
 226     // do not fail test just because a date could not be parsed
 227     if (verbose) {
 228       System.out.println("failed to parse date (assigning 'now') for: " + dateStr);
 229     }
 230     return null;
 231   }
 232
 233   @Override
 234   public void close() throws IOException {
 235     if (reader == null) {
 236       return;
 237     }
 238
 239     try {
 240       reader.close();
 241     } catch (IOException e) {
 242       if (verbose) {
 243         System.out.println("failed to close reader !");
 244         e.printStackTrace(System.out);
 245       }
 246     }
 247     reader = null;
 248   }
 249
 250   @Override
 251   public DocData getNextDocData(DocData docData) throws NoMoreDataException, IOException {
 252     String name = null;
 253     StringBuilder docBuf = getDocBuffer();
 254     ParsePathType parsedPathType;
 255
 256     // protect reading from the TREC files by multiple threads. The rest of the
 257     // method, i.e., parsing the content and returning the DocData can run unprotected.
 258     synchronized (lock) {
 259       if (reader == null) {
 260         openNextFile();
 261       }
 262
 263       // 1. skip until doc start - required for all TREC formats
 264       docBuf.setLength(0);
 265       read(docBuf, DOC, false, false);
 266
 267       // save parsedFile for passing trecDataParser after the sync block, in
 268       // case another thread will open another file in between.
 269       parsedPathType = currPathType;
 270
 271       // 2. name - required for all TREC formats
 272       docBuf.setLength(0);
 273       read(docBuf, DOCNO, true, false);
 274       name = docBuf.substring(DOCNO.length(), docBuf.indexOf(TERMINATING_DOCNO,
 275           DOCNO.length())).trim();
 276
 277       if (!excludeDocnameIteration) {
 278         name = name + "_" + iteration;
 279       }
 280
 281       // 3. read all until end of doc
 282       docBuf.setLength(0);
 283       read(docBuf, TERMINATING_DOC, false, true);
 284     }
 285
 286     // count char length of text to be parsed (may be larger than the resulted plain doc body text).
 287     addBytes(docBuf.length());
 288
 289     // This code segment relies on HtmlParser being thread safe. When we get
 290     // here, everything else is already private to that thread, so we're safe.
 291     try {
 292       docData = trecDocParser.parse(docData, name, this, docBuf, parsedPathType);
 293       addItem();
 294     } catch (InterruptedException ie) {
 295       throw new ThreadInterruptedException(ie);
 296     }
 297
 298     return docData;
 299   }
 300
 301   @Override
 302   public void resetInputs() throws IOException {
 303     synchronized (lock) {
 304       super.resetInputs();
 305       close();
 306       nextFile = 0;
 307       iteration = 0;
 308     }
 309   }
 310
 311   @Override
 312   public void setConfig(Config config) {
 313     super.setConfig(config);
 314     // dirs
 315     File workDir = new File(config.get("work.dir", "work"));
 316     String d = config.get("docs.dir", "trec");
 317     dataDir = new File(d);
 318     if (!dataDir.isAbsolute()) {
 319       dataDir = new File(workDir, d);
 320     }
 321     // files
 322     collectFiles(dataDir, inputFiles);
 323     if (inputFiles.size() == 0) {
 324       throw new IllegalArgumentException("No files in dataDir: " + dataDir);
 325     }
 326     // trec doc parser
 327     try {
 328       String trecDocParserClassName = config.get("trec.doc.parser", "org.apache.lucene.benchmark.byTask.feeds.TrecGov2Parser");
 329       trecDocParser = Class.forName(trecDocParserClassName).asSubclass(TrecDocParser.class).newInstance();
 330     } catch (Exception e) {
 331       // Should not get here. Throw runtime exception.
 332       throw new RuntimeException(e);
 333     }
 334     // html parser
 335     try {
 336       String htmlParserClassName = config.get("html.parser",
 337           "org.apache.lucene.benchmark.byTask.feeds.DemoHTMLParser");
 338       htmlParser = Class.forName(htmlParserClassName).asSubclass(HTMLParser.class).newInstance();
 339     } catch (Exception e) {
 340       // Should not get here. Throw runtime exception.
 341       throw new RuntimeException(e);
 342     }
 343     // encoding
 344     if (encoding == null) {
 345       encoding = "ISO-8859-1";
 346     }
 347     // iteration exclusion in doc name
 348     excludeDocnameIteration = config.get("content.source.excludeIteration", false);
 349   }
 350
 351 }