+++ /dev/null
-package org.apache.lucene.benchmark.byTask.feeds;
-
-/**
- * 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.
- */
-
-import java.io.BufferedReader;
-import java.io.File;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.InputStreamReader;
-import java.io.Reader;
-import java.text.DateFormat;
-import java.text.ParsePosition;
-import java.text.SimpleDateFormat;
-import java.util.ArrayList;
-import java.util.Date;
-import java.util.Locale;
-
-import org.apache.lucene.benchmark.byTask.feeds.TrecDocParser.ParsePathType;
-import org.apache.lucene.benchmark.byTask.utils.Config;
-import org.apache.lucene.benchmark.byTask.utils.StreamUtils;
-import org.apache.lucene.benchmark.byTask.utils.StringBuilderReader;
-import org.apache.lucene.util.ThreadInterruptedException;
-
-/**
- * Implements a {@link ContentSource} over the TREC collection.
- * <p>
- * Supports the following configuration parameters (on top of
- * {@link ContentSource}):
- * <ul>
- * <li><b>work.dir</b> - specifies the working directory. Required if "docs.dir"
- * denotes a relative path (<b>default=work</b>).
- * <li><b>docs.dir</b> - specifies the directory where the TREC files reside.
- * Can be set to a relative path if "work.dir" is also specified
- * (<b>default=trec</b>).
- * <li><b>trec.doc.parser</b> - specifies the {@link TrecDocParser} class to use for
- * parsing the TREC documents content (<b>default=TrecGov2Parser</b>).
- * <li><b>html.parser</b> - specifies the {@link HTMLParser} class to use for
- * parsing the HTML parts of the TREC documents content (<b>default=DemoHTMLParser</b>).
- * <li><b>content.source.encoding</b> - if not specified, ISO-8859-1 is used.
- * <li><b>content.source.excludeIteration</b> - if true, do not append iteration number to docname
- * </ul>
- */
-public class TrecContentSource extends ContentSource {
-
- private static final class DateFormatInfo {
- DateFormat[] dfs;
- ParsePosition pos;
- }
-
- public static final String DOCNO = "<DOCNO>";
- public static final String TERMINATING_DOCNO = "</DOCNO>";
- public static final String DOC = "<DOC>";
- public static final String TERMINATING_DOC = "</DOC>";
-
- /** separator between lines in the byffer */
- public static final String NEW_LINE = System.getProperty("line.separator");
-
- private static final String DATE_FORMATS [] = {
- "EEE, dd MMM yyyy kk:mm:ss z", // Tue, 09 Dec 2003 22:39:08 GMT
- "EEE MMM dd kk:mm:ss yyyy z", // Tue Dec 09 16:45:08 2003 EST
- "EEE, dd-MMM-':'y kk:mm:ss z", // Tue, 09 Dec 2003 22:39:08 GMT
- "EEE, dd-MMM-yyy kk:mm:ss z", // Tue, 09 Dec 2003 22:39:08 GMT
- "EEE MMM dd kk:mm:ss yyyy", // Tue Dec 09 16:45:08 2003
- "dd MMM yyyy", // 1 March 1994
- "MMM dd, yyyy", // February 3, 1994
- "yyMMdd", // 910513
- "hhmm z.z.z. MMM dd, yyyy", // 0901 u.t.c. April 28, 1994
- };
-
- private ThreadLocal<DateFormatInfo> dateFormats = new ThreadLocal<DateFormatInfo>();
- private ThreadLocal<StringBuilderReader> trecDocReader = new ThreadLocal<StringBuilderReader>();
- private ThreadLocal<StringBuilder> trecDocBuffer = new ThreadLocal<StringBuilder>();
- private File dataDir = null;
- private ArrayList<File> inputFiles = new ArrayList<File>();
- private int nextFile = 0;
- private int rawDocSize = 0;
-
- // Use to synchronize threads on reading from the TREC documents.
- private Object lock = new Object();
-
- // Required for test
- BufferedReader reader;
- int iteration = 0;
- HTMLParser htmlParser;
-
- private boolean excludeDocnameIteration;
- private TrecDocParser trecDocParser = new TrecGov2Parser(); // default
- ParsePathType currPathType; // not private for tests
-
- private DateFormatInfo getDateFormatInfo() {
- DateFormatInfo dfi = dateFormats.get();
- if (dfi == null) {
- dfi = new DateFormatInfo();
- dfi.dfs = new SimpleDateFormat[DATE_FORMATS.length];
- for (int i = 0; i < dfi.dfs.length; i++) {
- dfi.dfs[i] = new SimpleDateFormat(DATE_FORMATS[i], Locale.US);
- dfi.dfs[i].setLenient(true);
- }
- dfi.pos = new ParsePosition(0);
- dateFormats.set(dfi);
- }
- return dfi;
- }
-
- private StringBuilder getDocBuffer() {
- StringBuilder sb = trecDocBuffer.get();
- if (sb == null) {
- sb = new StringBuilder();
- trecDocBuffer.set(sb);
- }
- return sb;
- }
-
- Reader getTrecDocReader(StringBuilder docBuffer) {
- StringBuilderReader r = trecDocReader.get();
- if (r == null) {
- r = new StringBuilderReader(docBuffer);
- trecDocReader.set(r);
- } else {
- r.set(docBuffer);
- }
- return r;
- }
-
- HTMLParser getHtmlParser() {
- return htmlParser;
- }
-
- /**
- * Read until a line starting with the specified <code>lineStart</code>.
- * @param buf buffer for collecting the data if so specified/
- * @param lineStart line start to look for, must not be null.
- * @param collectMatchLine whether to collect the matching line into <code>buffer</code>.
- * @param collectAll whether to collect all lines into <code>buffer</code>.
- * @throws IOException
- * @throws NoMoreDataException
- */
- private void read(StringBuilder buf, String lineStart,
- boolean collectMatchLine, boolean collectAll) throws IOException, NoMoreDataException {
- String sep = "";
- while (true) {
- String line = reader.readLine();
-
- if (line == null) {
- openNextFile();
- continue;
- }
-
- rawDocSize += line.length();
-
- if (lineStart!=null && line.startsWith(lineStart)) {
- if (collectMatchLine) {
- buf.append(sep).append(line);
- sep = NEW_LINE;
- }
- return;
- }
-
- if (collectAll) {
- buf.append(sep).append(line);
- sep = NEW_LINE;
- }
- }
- }
-
- void openNextFile() throws NoMoreDataException, IOException {
- close();
- currPathType = null;
- while (true) {
- if (nextFile >= inputFiles.size()) {
- // exhausted files, start a new round, unless forever set to false.
- if (!forever) {
- throw new NoMoreDataException();
- }
- nextFile = 0;
- iteration++;
- }
- File f = inputFiles.get(nextFile++);
- if (verbose) {
- System.out.println("opening: " + f + " length: " + f.length());
- }
- try {
- InputStream inputStream = StreamUtils.inputStream(f); // support either gzip, bzip2, or regular text file, by extension
- reader = new BufferedReader(new InputStreamReader(inputStream, encoding), StreamUtils.BUFFER_SIZE);
- currPathType = TrecDocParser.pathType(f);
- return;
- } catch (Exception e) {
- if (verbose) {
- System.out.println("Skipping 'bad' file " + f.getAbsolutePath()+" due to "+e.getMessage());
- continue;
- }
- throw new NoMoreDataException();
- }
- }
- }
-
- public Date parseDate(String dateStr) {
- dateStr = dateStr.trim();
- DateFormatInfo dfi = getDateFormatInfo();
- for (int i = 0; i < dfi.dfs.length; i++) {
- DateFormat df = dfi.dfs[i];
- dfi.pos.setIndex(0);
- dfi.pos.setErrorIndex(-1);
- Date d = df.parse(dateStr, dfi.pos);
- if (d != null) {
- // Parse succeeded.
- return d;
- }
- }
- // do not fail test just because a date could not be parsed
- if (verbose) {
- System.out.println("failed to parse date (assigning 'now') for: " + dateStr);
- }
- return null;
- }
-
- @Override
- public void close() throws IOException {
- if (reader == null) {
- return;
- }
-
- try {
- reader.close();
- } catch (IOException e) {
- if (verbose) {
- System.out.println("failed to close reader !");
- e.printStackTrace(System.out);
- }
- }
- reader = null;
- }
-
- @Override
- public DocData getNextDocData(DocData docData) throws NoMoreDataException, IOException {
- String name = null;
- StringBuilder docBuf = getDocBuffer();
- ParsePathType parsedPathType;
-
- // protect reading from the TREC files by multiple threads. The rest of the
- // method, i.e., parsing the content and returning the DocData can run unprotected.
- synchronized (lock) {
- if (reader == null) {
- openNextFile();
- }
-
- // 1. skip until doc start - required for all TREC formats
- docBuf.setLength(0);
- read(docBuf, DOC, false, false);
-
- // save parsedFile for passing trecDataParser after the sync block, in
- // case another thread will open another file in between.
- parsedPathType = currPathType;
-
- // 2. name - required for all TREC formats
- docBuf.setLength(0);
- read(docBuf, DOCNO, true, false);
- name = docBuf.substring(DOCNO.length(), docBuf.indexOf(TERMINATING_DOCNO,
- DOCNO.length())).trim();
-
- if (!excludeDocnameIteration) {
- name = name + "_" + iteration;
- }
-
- // 3. read all until end of doc
- docBuf.setLength(0);
- read(docBuf, TERMINATING_DOC, false, true);
- }
-
- // count char length of text to be parsed (may be larger than the resulted plain doc body text).
- addBytes(docBuf.length());
-
- // This code segment relies on HtmlParser being thread safe. When we get
- // here, everything else is already private to that thread, so we're safe.
- try {
- docData = trecDocParser.parse(docData, name, this, docBuf, parsedPathType);
- addDoc();
- } catch (InterruptedException ie) {
- throw new ThreadInterruptedException(ie);
- }
-
- return docData;
- }
-
- @Override
- public void resetInputs() throws IOException {
- synchronized (lock) {
- super.resetInputs();
- close();
- nextFile = 0;
- iteration = 0;
- }
- }
-
- @Override
- public void setConfig(Config config) {
- super.setConfig(config);
- // dirs
- File workDir = new File(config.get("work.dir", "work"));
- String d = config.get("docs.dir", "trec");
- dataDir = new File(d);
- if (!dataDir.isAbsolute()) {
- dataDir = new File(workDir, d);
- }
- // files
- collectFiles(dataDir, inputFiles);
- if (inputFiles.size() == 0) {
- throw new IllegalArgumentException("No files in dataDir: " + dataDir);
- }
- // trec doc parser
- try {
- String trecDocParserClassName = config.get("trec.doc.parser", "org.apache.lucene.benchmark.byTask.feeds.TrecGov2Parser");
- trecDocParser = Class.forName(trecDocParserClassName).asSubclass(TrecDocParser.class).newInstance();
- } catch (Exception e) {
- // Should not get here. Throw runtime exception.
- throw new RuntimeException(e);
- }
- // html parser
- try {
- String htmlParserClassName = config.get("html.parser",
- "org.apache.lucene.benchmark.byTask.feeds.DemoHTMLParser");
- htmlParser = Class.forName(htmlParserClassName).asSubclass(HTMLParser.class).newInstance();
- } catch (Exception e) {
- // Should not get here. Throw runtime exception.
- throw new RuntimeException(e);
- }
- // encoding
- if (encoding == null) {
- encoding = "ISO-8859-1";
- }
- // iteration exclusion in doc name
- excludeDocnameIteration = config.get("content.source.excludeIteration", false);
- }
-
-}
fnp / pylucene.git / blobdiff