1 package org.apache.lucene.benchmark.byTask.feeds;
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
11 * http://www.apache.org/licenses/LICENSE-2.0
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.
20 import java.io.BufferedReader;
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;
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;
40 * Implements a {@link ContentSource} over the TREC collection.
42 * Supports the following configuration parameters (on top of
43 * {@link ContentSource}):
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
58 public class TrecContentSource extends ContentSource {
60 private static final class DateFormatInfo {
65 public static final String DOCNO = "<DOCNO>";
66 public static final String TERMINATING_DOCNO = "</DOCNO>";
67 public static final String DOC = "<DOC>";
68 public static final String TERMINATING_DOC = "</DOC>";
70 /** separator between lines in the byffer */
71 public static final String NEW_LINE = System.getProperty("line.separator");
73 private static final String DATE_FORMATS [] = {
74 "EEE, dd MMM yyyy kk:mm:ss z", // Tue, 09 Dec 2003 22:39:08 GMT
75 "EEE MMM dd kk:mm:ss yyyy z", // Tue Dec 09 16:45:08 2003 EST
76 "EEE, dd-MMM-':'y kk:mm:ss z", // Tue, 09 Dec 2003 22:39:08 GMT
77 "EEE, dd-MMM-yyy kk:mm:ss z", // Tue, 09 Dec 2003 22:39:08 GMT
78 "EEE MMM dd kk:mm:ss yyyy", // Tue Dec 09 16:45:08 2003
79 "dd MMM yyyy", // 1 March 1994
80 "MMM dd, yyyy", // February 3, 1994
82 "hhmm z.z.z. MMM dd, yyyy", // 0901 u.t.c. April 28, 1994
85 private ThreadLocal<DateFormatInfo> dateFormats = new ThreadLocal<DateFormatInfo>();
86 private ThreadLocal<StringBuilderReader> trecDocReader = new ThreadLocal<StringBuilderReader>();
87 private ThreadLocal<StringBuilder> trecDocBuffer = new ThreadLocal<StringBuilder>();
88 private File dataDir = null;
89 private ArrayList<File> inputFiles = new ArrayList<File>();
90 private int nextFile = 0;
91 private int rawDocSize = 0;
93 // Use to synchronize threads on reading from the TREC documents.
94 private Object lock = new Object();
97 BufferedReader reader;
99 HTMLParser htmlParser;
101 private boolean excludeDocnameIteration;
102 private TrecDocParser trecDocParser = new TrecGov2Parser(); // default
103 ParsePathType currPathType; // not private for tests
105 private DateFormatInfo getDateFormatInfo() {
106 DateFormatInfo dfi = dateFormats.get();
108 dfi = new DateFormatInfo();
109 dfi.dfs = new SimpleDateFormat[DATE_FORMATS.length];
110 for (int i = 0; i < dfi.dfs.length; i++) {
111 dfi.dfs[i] = new SimpleDateFormat(DATE_FORMATS[i], Locale.US);
112 dfi.dfs[i].setLenient(true);
114 dfi.pos = new ParsePosition(0);
115 dateFormats.set(dfi);
120 private StringBuilder getDocBuffer() {
121 StringBuilder sb = trecDocBuffer.get();
123 sb = new StringBuilder();
124 trecDocBuffer.set(sb);
129 Reader getTrecDocReader(StringBuilder docBuffer) {
130 StringBuilderReader r = trecDocReader.get();
132 r = new StringBuilderReader(docBuffer);
133 trecDocReader.set(r);
140 HTMLParser getHtmlParser() {
145 * Read until a line starting with the specified <code>lineStart</code>.
146 * @param buf buffer for collecting the data if so specified/
147 * @param lineStart line start to look for, must not be null.
148 * @param collectMatchLine whether to collect the matching line into <code>buffer</code>.
149 * @param collectAll whether to collect all lines into <code>buffer</code>.
150 * @throws IOException
151 * @throws NoMoreDataException
153 private void read(StringBuilder buf, String lineStart,
154 boolean collectMatchLine, boolean collectAll) throws IOException, NoMoreDataException {
157 String line = reader.readLine();
164 rawDocSize += line.length();
166 if (lineStart!=null && line.startsWith(lineStart)) {
167 if (collectMatchLine) {
168 buf.append(sep).append(line);
175 buf.append(sep).append(line);
181 void openNextFile() throws NoMoreDataException, IOException {
185 if (nextFile >= inputFiles.size()) {
186 // exhausted files, start a new round, unless forever set to false.
188 throw new NoMoreDataException();
193 File f = inputFiles.get(nextFile++);
195 System.out.println("opening: " + f + " length: " + f.length());
198 InputStream inputStream = StreamUtils.inputStream(f); // support either gzip, bzip2, or regular text file, by extension
199 reader = new BufferedReader(new InputStreamReader(inputStream, encoding), StreamUtils.BUFFER_SIZE);
200 currPathType = TrecDocParser.pathType(f);
202 } catch (Exception e) {
204 System.out.println("Skipping 'bad' file " + f.getAbsolutePath()+" due to "+e.getMessage());
207 throw new NoMoreDataException();
212 public Date parseDate(String dateStr) {
213 dateStr = dateStr.trim();
214 DateFormatInfo dfi = getDateFormatInfo();
215 for (int i = 0; i < dfi.dfs.length; i++) {
216 DateFormat df = dfi.dfs[i];
218 dfi.pos.setErrorIndex(-1);
219 Date d = df.parse(dateStr, dfi.pos);
225 // do not fail test just because a date could not be parsed
227 System.out.println("failed to parse date (assigning 'now') for: " + dateStr);
233 public void close() throws IOException {
234 if (reader == null) {
240 } catch (IOException e) {
242 System.out.println("failed to close reader !");
243 e.printStackTrace(System.out);
250 public DocData getNextDocData(DocData docData) throws NoMoreDataException, IOException {
252 StringBuilder docBuf = getDocBuffer();
253 ParsePathType parsedPathType;
255 // protect reading from the TREC files by multiple threads. The rest of the
256 // method, i.e., parsing the content and returning the DocData can run unprotected.
257 synchronized (lock) {
258 if (reader == null) {
262 // 1. skip until doc start - required for all TREC formats
264 read(docBuf, DOC, false, false);
266 // save parsedFile for passing trecDataParser after the sync block, in
267 // case another thread will open another file in between.
268 parsedPathType = currPathType;
270 // 2. name - required for all TREC formats
272 read(docBuf, DOCNO, true, false);
273 name = docBuf.substring(DOCNO.length(), docBuf.indexOf(TERMINATING_DOCNO,
274 DOCNO.length())).trim();
276 if (!excludeDocnameIteration) {
277 name = name + "_" + iteration;
280 // 3. read all until end of doc
282 read(docBuf, TERMINATING_DOC, false, true);
285 // count char length of text to be parsed (may be larger than the resulted plain doc body text).
286 addBytes(docBuf.length());
288 // This code segment relies on HtmlParser being thread safe. When we get
289 // here, everything else is already private to that thread, so we're safe.
291 docData = trecDocParser.parse(docData, name, this, docBuf, parsedPathType);
293 } catch (InterruptedException ie) {
294 throw new ThreadInterruptedException(ie);
301 public void resetInputs() throws IOException {
302 synchronized (lock) {
311 public void setConfig(Config config) {
312 super.setConfig(config);
314 File workDir = new File(config.get("work.dir", "work"));
315 String d = config.get("docs.dir", "trec");
316 dataDir = new File(d);
317 if (!dataDir.isAbsolute()) {
318 dataDir = new File(workDir, d);
321 collectFiles(dataDir, inputFiles);
322 if (inputFiles.size() == 0) {
323 throw new IllegalArgumentException("No files in dataDir: " + dataDir);
327 String trecDocParserClassName = config.get("trec.doc.parser", "org.apache.lucene.benchmark.byTask.feeds.TrecGov2Parser");
328 trecDocParser = Class.forName(trecDocParserClassName).asSubclass(TrecDocParser.class).newInstance();
329 } catch (Exception e) {
330 // Should not get here. Throw runtime exception.
331 throw new RuntimeException(e);
335 String htmlParserClassName = config.get("html.parser",
336 "org.apache.lucene.benchmark.byTask.feeds.DemoHTMLParser");
337 htmlParser = Class.forName(htmlParserClassName).asSubclass(HTMLParser.class).newInstance();
338 } catch (Exception e) {
339 // Should not get here. Throw runtime exception.
340 throw new RuntimeException(e);
343 if (encoding == null) {
344 encoding = "ISO-8859-1";
346 // iteration exclusion in doc name
347 excludeDocnameIteration = config.get("content.source.excludeIteration", false);